[Paper Review] Towards Green AI in Fine-tuning Large Language Models via Adaptive Backpropagation

Kai Huang, Hanyun Yin+2 more|arXiv (Cornell University)|Sep 22, 2023

Tracheal and airway disordersMedicine3 citations

TL;DR

GreenTrainer is a novel fine-tuning method for large language models that adaptively selects trainable tensors based on their backpropagation cost and contribution to model accuracy, reducing FLOPs by up to 64% without significant accuracy loss. It outperforms LoRA by achieving up to 4% higher accuracy at comparable FLOPs reduction, enabling greener AI training through dynamic, objective-driven computation optimization.

ABSTRACT

Fine-tuning is the most effective way of adapting pre-trained large language models (LLMs) to downstream applications. With the fast growth of LLM-enabled AI applications and democratization of open-souced LLMs, fine-tuning has become possible for non-expert individuals, but intensively performed LLM fine-tuning worldwide could result in significantly high energy consumption and carbon footprint, which may bring large environmental impact. Mitigating such environmental impact towards Green AI directly correlates to reducing the FLOPs of fine-tuning, but existing techniques on efficient LLM fine-tuning can only achieve limited reduction of such FLOPs, due to their ignorance of the backpropagation cost in fine-tuning. To address this limitation, in this paper we present GreenTrainer, a new LLM fine-tuning technique that adaptively evaluates different tensors' backpropagation costs and contributions to the fine-tuned model accuracy, to minimize the fine-tuning cost by selecting the most appropriate set of tensors in training. Such selection in GreenTrainer is made based on a given objective of FLOPs reduction, which can flexibly adapt to the carbon footprint in energy supply and the need in Green AI. Experiment results over multiple open-sourced LLM models and abstractive summarization datasets show that, compared to fine-tuning the whole LLM model, GreenTrainer can save up to 64% FLOPs in fine-tuning without any noticeable model accuracy loss. Compared to the existing fine-tuning techniques such as LoRa, GreenTrainer can achieve up to 4% improvement on model accuracy with on-par FLOPs reduction.

Motivation & Objective

• To address the high energy consumption and carbon footprint of widespread LLM fine-tuning by reducing FLOPs in training.
• To overcome the limitation of existing methods that focus only on weight update FLOPs and ignore backpropagation costs.
• To develop a technique that adaptively selects the most impactful tensors for backpropagation based on a flexible FLOPs reduction objective.
• To minimize accuracy loss during FLOPs reduction by dynamically adjusting the selection of trainable tensors based on energy supply carbon footprint.
• To enable sustainable, green AI fine-tuning that balances computational efficiency and model performance across diverse energy environments.

Proposed method

• GreenTrainer adaptively evaluates the importance of individual neural network tensors in backpropagation based on their contribution to model accuracy and computational cost.
• It formulates a flexible FLOPs reduction objective that can be adjusted according to the carbon footprint of energy supply, enabling dynamic trade-offs between efficiency and accuracy.
• The method selectively enables backpropagation only for the most critical tensors, reducing FLOPs in activation gradient computation, which accounts for at least 33% of total training FLOPs.
• It uses a runtime tensor importance metric that avoids reliance on input attribution methods, which are inapplicable to model parameters.
• The approach integrates with existing LoRA-style parameter-efficient fine-tuning by focusing on selective backpropagation rather than parameter updates.
• It supports various LLM sizes and architectures, including OPT and LLaMA, and is open-sourced for community use.

Experimental results

Research questions

• RQ1 Can selective backpropagation significantly reduce FLOPs in LLM fine-tuning while preserving model accuracy?
• RQ2 How does the adaptive selection of trainable tensors based on backpropagation cost and contribution affect fine-tuning efficiency and accuracy?
• RQ3 To what extent can GreenTrainer reduce FLOPs compared to full fine-tuning and existing methods like LoRA?
• RQ4 How does the performance of GreenTrainer vary across different LLM sizes and downstream tasks?
• RQ5 Can the method be flexibly tuned to align with varying carbon footprint levels in energy supply?

Key findings

• GreenTrainer reduces fine-tuning FLOPs by up to 64% compared to full fine-tuning, with no noticeable accuracy loss across multiple datasets and models.
• On the OPT-6.7B model, GT-0.4 saves 60% more FLOPs than LoRA on the SciTLDR dataset while maintaining the same model accuracy.
• Compared to LoRA, GreenTrainer achieves up to 4% higher ROUGE scores on the SciTLDR dataset at the same level of FLOPs reduction.
• For the OPT-350M model, GT-0.5 saves 50% of FLOPs and achieves 2%-4% higher accuracy than LoRA on the SciTDR dataset.
• On larger models like OPT-2.7B and OPT-6.7B, GreenTrainer consistently outperforms LoRA in FLOPs reduction and maintains or improves accuracy.
• The method scales effectively across model sizes, from 350M to 6.7B parameters, showing consistent advantages in both efficiency and performance.