[Paper Review] Sparse Fine-tuning for Inference Acceleration of Large Language Models

Eldar Kurtic, Denis Kuznedelev+3 more|arXiv (Cornell University)|Oct 10, 2023

Topic Modeling4 citations

TL;DR

This paper proposes a sparse fine-tuning framework for large language models (LLMs) that combines a novel L2-based knowledge distillation method, SquareHead, with iterative pruning to maintain high accuracy even at extreme sparsity levels (up to 75%). The approach enables significant end-to-end inference speedups—up to 9.08x on CPU and 7.49x on GPU—while preserving accuracy, and demonstrates compatibility with quantization for further performance gains on both CPU and GPU workloads.

ABSTRACT

We consider the problem of accurate sparse fine-tuning of large language models (LLMs), that is, fine-tuning pretrained LLMs on specialized tasks, while inducing sparsity in their weights. On the accuracy side, we observe that standard loss-based fine-tuning may fail to recover accuracy, especially at high sparsities. To address this, we perform a detailed study of distillation-type losses, determining an L2-based distillation approach we term SquareHead which enables accurate recovery even at higher sparsities, across all model types. On the practical efficiency side, we show that sparse LLMs can be executed with speedups by taking advantage of sparsity, for both CPU and GPU runtimes. While the standard approach is to leverage sparsity for computational reduction, we observe that in the case of memory-bound LLMs sparsity can also be leveraged for reducing memory bandwidth. We exhibit end-to-end results showing speedups due to sparsity, while recovering accuracy, on T5 (language translation), Whisper (speech translation), and open GPT-type (MPT for text generation). For MPT text generation, we show for the first time that sparse fine-tuning can reach 75% sparsity without accuracy drops, provide notable end-to-end speedups for both CPU and GPU inference, and highlight that sparsity is also compatible with quantization approaches. Models and software for reproducing our results are provided in Section 6.

Motivation & Objective

• Address the challenge of maintaining model accuracy during sparse fine-tuning of large language models, especially at high sparsity levels.
• Overcome training instability and poor accuracy recovery in sparse fine-tuning due to limited fine-tuning data and model complexity.
• Demonstrate that sparsity can be effectively leveraged not only for computational reduction but also for reducing memory bandwidth in memory-bound LLM inference.
• Show that sparsity is compatible with quantization, enabling compounded speedups on both CPU and GPU hardware.
• Provide reproducible code and models to enable adoption and extension of the method across diverse LLM applications.

Proposed method

• Apply iterative pruning during fine-tuning, gradually increasing sparsity while retraining the model on downstream tasks.
• Introduce SquareHead, a per-token ℓ₂-based knowledge distillation loss that distills both output logits and intermediate hidden representations from a dense teacher model to a sparse student model.
• Use a distillation loss that minimizes the L2 difference between teacher and student activations at each layer, improving generalization and stability during sparse fine-tuning.
• Leverage the DeepSparse inference runtime to exploit sparsity for CPU inference speedups via sparse GEMM kernels.
• Implement GPU-aware N:M sparse formats to enable efficient sparse inference on GPUs, combining sparsity with quantization (INT8).
• Apply post-training quantization to sparse models and evaluate the combined effect of sparsity and quantization on accuracy and latency.

Figure 1 : Speedups for sparse CPU and GPU kernels on an MPT-7B layer.

Experimental results

Research questions

• RQ1 Can sparse fine-tuning maintain high accuracy in large language models when fine-tuning data is limited and sparsity is high?
• RQ2 How does knowledge distillation, particularly SquareHead, compare to standard cross-entropy and KL-based distillation in preserving accuracy during sparse fine-tuning?
• RQ3 To what extent can sparsity reduce memory bandwidth pressure in memory-bound LLM inference workloads?
• RQ4 Can sparsity and quantization be jointly applied to achieve greater inference speedups without significant accuracy degradation?
• RQ5 What are the end-to-end speedup and accuracy trade-offs of sparse fine-tuning across diverse LLM applications such as translation, speech transcription, and reasoning?

Key findings

• SquareHead distillation enables accurate recovery of model performance even at 75% sparsity, with the sparse MPT-7B model achieving 28.6% test accuracy on GSM8K—only a 1.6% drop from the dense FP32 baseline of 28.2%.
• The 60% sparse INT8 MPT-7B model achieves a 6.7x speedup on a single-core AMD Ryzen 9 7950X, and the 80% sparse INT8 model reaches 9.08x speedup with a 7% accuracy drop.
• On CPU, the 70% sparse MPT-7B model achieves 7.49x speedup at 21.1 tokens/second on 4 cores, demonstrating that sparsity alone enables substantial gains.
• The sparse fine-tuned MPT-7B model with SquareHead KD outperforms the dense baseline, achieving 33.0% accuracy on GSM8K, indicating that distillation improves generalization.
• Sparsity reduces memory bandwidth pressure in memory-bound generative models, enabling speedups even on CPU systems where computation is not the bottleneck.
• The method is compatible with both FP32 and INT8 quantization, and the combined use of sparsity and quantization leads to compounded performance improvements on both CPU and GPU.

Figure 2 : BLEU ( $\uparrow$ ) score for variants of loss functions on English-German WMT14 and T5-Small at various compression rates.