Towards Green AI in Fine-tuning Large Language Models via Adaptive Backpropagation

Dear reviewer,

Thank you so much for your efforts on reviewing our paper. Please find our responses to your comments and questions below.

The template is highly altered to fit in more content. For example, the author section is removed, which may violate the policy.

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.

The paper misses some references to existing work that uses similar techniques ...

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.

There are some other efficient training methods focusing more on FLOPs reduction instead of parameter reduction like [d] ...

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.

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

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.

The tensor importance estimation is based on the gradient information. Is it consistent throughout the training, or will the importance change?

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.

Dear Reviewer,

We are writing to kindly follow up with our previous response to your review comments, and would like to kindly enquire whether our response fully answered your questions and addressed your concerns to our paper. We will be more than happy to provide further details and responses for any additional questions you may have. Again, thank you so much for your review efforts and your insightful feedback about our paper!

Dear Reviewer,

Thank you so much for your response, and we are glad to see that most of your questions and concerns have been addressed. Regarding your question about using a fixed set of selected tensors determined at the initial stage of training, please find our further clarifications as below.

We would like to clarify that the correlation between the trained model's accuracy and frequency of tensor importance evaluation (and tensor selection) mainly depends on how much the tensors' importances will vary. In our previous responses to other reviewers, we justified that such variance in one epoch would be very small, and so that it is unnecessary to do multiple times of tensor importance evaluation and selection within one epoch. The possible variations of tensor importances over different epochs, on the other hand, is also related to the scenario of LLM fine-tuning and the corresponding characteristics of training dataset.

We agree with the reviewer that, if the LLM fine-tuning uses a fixed training dataset, it is possible that using a fixed tensor selection decided at the initial phase of training may not result in a significant model accuracy drop, compared to runtime tensor selection. However, in practical LLM fine-tuning scenarios, this assumption usually does not hold due to the following two reasons. First, a lot of LLM fine-tuning scenarios, such as online learning and model personalization, is continuously training the model using online data, which is continuously being generated at runtime with varying data distributions. Such varying data distributions in training will surely result in different importances of tensors through the training procedure and hence require runtime tensor selection. Such online LLM fine-tuning scenarios is recently getting more and more popular, especially with the possibility of deploying LLMs onto user's personal mobile devices such as smartphones. Second, even for a fixed training dataset, it is also possible that the importances of some tensors may change as the training progresses. In these cases, dynamic tensor selection could improve the trained model accuracy. To verify this, we conducted additional experiments using OPT-2.7B model on the webquestions dataset and generative QA task. The results are as follows:

Note that, such improvement of model accuracy would be dataset and model dependent, but these experiment results above demonstrated the necessity of runtime tensor selection. In addition, our experiment results also showed that such tensor importance evaluation and selection indeed incur very little extra computing overhead.

We hope that these clarifications fully addressed your concerns, and we would be more than happy to discuss further. Thank you so much!

Dear reviewer,

Thank you so much for your efforts on reviewing our paper. Please find our responses to your comments and questions below.

In the experiments, the authors only fine-tune the model for 5 epochs ....

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:

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.

How the running time is measured? There is a lack of hardware specifications.

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.

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.

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

The writing of the paper is sometimes confusing ...

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.

The proposed dynamic programming can only find approximate solutions ....

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.

Dear Reviewer,

We are writing to kindly follow up with our previous response to your review comments, and would like to kindly enquire whether our response fully answered your questions and addressed your concerns to our paper. We will be more than happy to provide further details and responses for any additional questions you may have. Again, thank you so much for your review efforts and your insightful feedback about our paper!

Dear reviewer,

Thank you so much for your efforts on reviewing our paper. Please find our responses to your comments and questions below. We hope these responses fully address your concerns on our paper, and if so we would highly appreciate if you could reconsider your rating of our paper. Please feel free to let us know if you would like further details about any part of the paper, and we would be more than happy to discuss further. Thank you!

What is the e2e finetuning time of GreenTrainer compared to baselines.

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.

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?

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.

What is the m binary vector, or at least the number of important tensors, for the experimental results?

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.

Dear reviewer,

Thank you so much for your efforts on reviewing our paper. Please find our responses to your comments and questions below.

It is unclear whether the fine-tuning process includes the DP (data parallelism) overhead ....

We are not sure whether the reviewer refers DP to dynamic programming (Section 3.3) or data parallelism.

For dynamic programming: all the experiments we did in this paper were evaluating the end-to-end fine-tuning costs, including the extra computations and time needed for both tensor importance evaluation and DP. In Appendix A.2, we provided experiment results of the DP overhead in fine-tuning, which can be flexibly adjusted by controlling its resolution. Our experiment results in Table 7 showed that, by selecting the appropriate DP resolution, we can effectively restrain the DP overhead to be <1% of the entire fine-tuning cost.

For data parallelism: we did our experiments on a single H100 80GB GPU and hence did not involve any communication cost between multiple GPUs for data parallelism.

How X% of FLOPs saving can lead to a X% reduction in latency in Tab 2? ....

We agree that the amount of FLOPs saving can be different from wall-clock time reduction, due to memory access overhead. In our case, the percentage of such memory access overhead is small due to the strong hardware (H100 80GB GPU), but still noticeable. In Table 1, the percentage of wall-clock time saving in GT-0.4/0.5/0.7 is slightly lower than FLOPs saving, indicating the extra cost of memory access due to tensor importance evaluation and DP. Such difference grows as $\rho$ increases, but is within 5%. In contrast, the memory access overhead in the baselines (e.g., FT-Top2, Prefix-T and LoRA) is generally lower. Even if we include memory access overhead, GT can still achieve higher wall-clock time savings.

The reviewer also mentioned an example that updating the last 2 blocks has similar latency with updating a partial set of all blocks. This is because in OPT and BLOOMZ models the output embedding layer is tied with the input embedding layer. As long as the output layer is selected to be trained, the input layer will be backpropagated for updating, leading to a maximum backpropagation path. In contrast, FLAN-T5 doesn’t tie input and output embeddings, and so FT-Top2’s amount of savings on FLAT-T5 is higher.

More experiments should be conducted to compare it against baselines ....

We conducted experiments using OPT-2.7B, on webquestions and piqa datasets for generative QA tasks. We evaluate the sentence-level accuracy which requires the generated answer to exactly match the ground truth. Our results are as follows:

LLMs are typically over-parameterized, and updating a random set of sparse parameters can achieve similar performance. Baselines such as random selection, Fisher, and SNIP should be discussed.

We agree that LLMs are typically over-parametrized, but selecting random parameter sets usually leads to much slower convergence speed, and we need more training epochs in fine-tuning. In addition, since the correlation between an arbitrarily selected set of parameters and the LLM’s behavior is unknown, it could be highly random and time consuming to find a best set of parameters that maximizes accuracy.

We provided experiment results in Table 4 about 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$). SNIP ($|\partial L/\partial w * w|$) and Fisher information ($|\partial L/\partial w \cdot w|^2$) are used in pruning to quantify the weight importance at current values. In contrast, our tensor importance evaluation is to quantify the importance of performing weight updates to the reduction of training loss ($-\partial L/\partial w \cdot \Delta w$). In other words, we measure the importance of “tensor update” ($\Delta w$) instead of the importance of the tensor’s current value ($w$).

Furthermore, our metric is directly derived from the first-order approximation after Taylor expansion of the training loss. Hence, it retains a moderate level of additivity in low learning rate ($10^-5$) regimes for fine-tuning.

There are concerns that this overhead of DP will grow with the size of the model and soon become non-negligible ...

In appendix A.2, we provided experiment results about the computing cost of our proposed DP approach, and showed that, by adjusting the resolution of DP, we can restrain the computing cost of DP within 1% of the entire fine-tuning cost, without affecting accuracy.

Dear Reviewer,

We are writing to kindly follow up with our previous response to your review comments, and would like to kindly enquire whether our response fully answered your questions and addressed your concerns to our paper. We will be more than happy to provide further details and responses for any additional questions you may have. Again, thank you so much for your review efforts and your insightful feedback about our paper!

After reading the authors' reponse (to me and other reviewers), some of my questions are addressed. However, I am showing concerns

• ** restrain the DP overhead to be <1% of the entire fine-tuning cost**

As mentioned in the reponse to #P1Bv, the DP is performed per epoch and this overhead makes sense. However, since the DP search is only performed once per epoch and the gradient / weight may vary a lot during fine-tuning. Does this suggest that GreenTrainer cannot handle large dataset (there are a lot of iterations in one epoch).

• the percentage of such memory access overhead is small due to the strong hardware (H100 80GB GPU)

Does this claim limit the applibility of "Green" Trainer? Means that the training can be green and cheap only when you have fancy H100 hardware?

• We conducted experiments using OPT-2.7B, on webquestions and piqa datasets for generative QA tasks

This PIQA accuracy looks weird. According to original OPT paper, Appendix A, the lowest performance of OPT is near 62.5 but here the LoRA accuracy is only 49.5%, which is a huge drop.