PAT: Pruning-Aware Tuning for Large Language Models

TL;DR

This paper introduces Pruning-Aware Tuning (PAT), a method that combines structural pruning with fine-tuning of large language models to reduce redundancy while maintaining high performance, leading to faster and more efficient models.

Contribution

The paper proposes a novel PAT paradigm with Hybrid Sparsification Modules and Identity Loss to effectively integrate pruning into fine-tuning, improving efficiency without sacrificing accuracy.

Findings

Achieves 1.33× speedup with 25% pruning on Llama2-7b.

Outperforms LoRA fine-tuning by up to 1.26% in accuracy.

Demonstrates effectiveness across large language models.

Abstract

Large language models (LLMs) excel in language tasks, especially with supervised fine-tuning after pre-training. However, their substantial memory and computational requirements hinder practical applications. Structural pruning, which reduces less significant weight dimensions, is one solution. Yet, traditional post-hoc pruning often leads to significant performance loss, with limited recovery from further fine-tuning due to reduced capacity. Since the model fine-tuning refines the general and chaotic knowledge in pre-trained models, we aim to incorporate structural pruning with the fine-tuning, and propose the Pruning-Aware Tuning (PAT) paradigm to eliminate model redundancy while preserving the model performance to the maximum extend. Specifically, we insert the innovative Hybrid Sparsification Modules (HSMs) between the Attention and FFN components to accordingly sparsify the upstream and downstream linear modules. The HSM comprises a lightweight operator and a globally shared trainable mask. The lightweight operator maintains a training overhead comparable to that of LoRA, while the trainable mask unifies the channels to be sparsified, ensuring structural pruning. Additionally, we propose the Identity Loss which decouples the transformation and scaling properties of the HSMs to enhance training robustness. Extensive experiments demonstrate that PAT excels in both performance and efficiency. For example, our Llama2-7b model with a 25\% pruning ratio achieves 1.33$\times$ speedup while outperforming the LoRA-finetuned model by up to 1.26\% in accuracy with a similar training cost. Code: https://github.com/kriskrisliu/PAT_Pruning-Aware-Tuning