All-in-One Tuning and Structural Pruning for Domain-Specific LLMs

TL;DR

This paper introduces ATP, a unified approach that combines pruning and fine-tuning of large language models in one stage, dynamically optimizing substructures during training for better domain-specific performance.

Contribution

The paper proposes a novel one-stage pruning and fine-tuning method that adapts pruning decisions during training, improving domain-specific LLM performance over traditional two-stage methods.

Findings

ATP recovers up to 88% of dense model performance after 40% pruning.

ATP outperforms state-of-the-art two-stage pruning methods in legal and healthcare tasks.

ATP effectively integrates LoRA-aware regularization for better pruning during low-data fine-tuning.

Abstract

Existing pruning techniques for large language models (LLMs) targeting domain-specific applications typically follow a two-stage process: pruning the pretrained general-purpose LLMs and then fine-tuning the pruned LLMs on specific domains. However, the pruning decisions, derived from the pretrained weights, remain unchanged during fine-tuning, even if the weights have been updated. Therefore, such a combination of the pruning decisions and the finetuned weights may be suboptimal, leading to non-negligible performance degradation. To address these limitations, we propose ATP: All-in-One Tuning and Structural Pruning, a unified one-stage structural pruning and fine-tuning approach that dynamically identifies the current optimal substructure throughout the fine-tuning phase via a trainable pruning decision generator. Moreover, given the limited available data for domain-specific applications, Low-Rank Adaptation (LoRA) becomes a common technique to fine-tune the LLMs. In ATP, we introduce LoRA-aware forward and sparsity regularization to ensure that the substructures corresponding to the learned pruning decisions can be directly removed after the ATP process. ATP outperforms the state-of-the-art two-stage pruning methods on tasks in the legal and healthcare domains. More specifically, ATP recovers up to 88% and 91% performance of the dense model when pruning 40% parameters of LLaMA2-7B and LLaMA3-8B models, respectively.