Sparse Fine-tuning for Inference Acceleration of Large Language Models

Sparse Fine-tuning for Accelerated LLM Inference

The paper "Sparse Fine-tuning for Inference Acceleration of LLMs" by Eldar Kurtic et al. explore the challenge of efficiently fine-tuning LLMs while imposing sparsity on their weights. This research endeavor addresses two primary objectives: achieving high accuracy in sparse LLMs and realizing practical efficiency during inference through sparsity-based acceleration.

Key Contributions and Findings

1. Accuracy Recovery with High Sparsity: The paper identifies that conventional fine-tuning methods often fail to maintain accuracy at high sparsity levels. To mitigate this issue, the authors propose a distillation-based loss function named SquareHead, which incorporates L2-based distillation. Experimental results demonstrate that SquareHead enables effective accuracy recovery across various LLM types, even at higher sparsity levels.

2. Efficiency Gains: On the efficiency front, the authors show that leveraging sparsity can significantly speed up LLM inference on both CPUs and GPUs. They demonstrate this through the execution of sparse LLMs, such as T5 for language translation, Whisper for speech translation, and open GPT-type models like MPT for text generation. For instance, the MPT model achieved up to 75% sparsity without any accuracy decline. Moreover, they highlight that sparsity can be effectively combined with quantization techniques, further enhancing the computational efficiency.

3. Memory-Bound Inference: The research reveals a distinctive benefit of sparsity in mitigating memory bandwidth constraints, particularly for memory-bound LLM inference tasks. The study provides empirical evidence that compressed sparse weights, which are decompressed on-the-fly during computation, can lead to substantial speedups.

Methodology

The methodology consists of several components designed to achieve sparse fine-tuning while maintaining high accuracy:

Sparsification:

Sparse LLMs are obtained by incrementally imposing higher sparsity levels during fine-tuning. This iterative approach ensures that the model adapts to sparsity progressively, thereby reducing the risk of instability and divergence.

Distillation Strategies:

The research evaluates various distillation strategies, including standard knowledge distillation (KD) using cross-entropy loss, and the proposed SquareHead KD, which incorporates intermediate layer distillation using normalized mean squared error (MSE) loss. SquareHead consistently outperforms other strategies by recovering accuracy even at higher sparsities.

Runtime Acceleration:

The runtime benefits of sparsity are realized through sophisticated algorithmic optimizations in sparse matrix computations. For GPUs, custom CUDA kernels are employed for efficient execution of N:M sparsity patterns, achieving notable speedups over dense execution.

Numerical Results and Practical Implications

The empirical evaluation spans various models and tasks, yielding robust findings:

• For T5 model fine-tuning on the WMT14 English-German translation task, sparsity levels up to 75% are achieved with minimal performance degradation (BLEU scores). This translates to over 2x inference speedup on CPUs.
• The Whisper model for Hindi ASR shows competitive Word Error Rates (WER) at 70-80% sparsity, with speedups reaching 2.5x on CPUs.
• MPT-7B generative models fine-tuned on the GSM8K dataset successfully maintain performance at 70% sparsity, with a substantial speedup in generating tokens per second. Moreover, combined sparsity and INT8 quantization leads to remarkable gains in decoding speed, exemplified by a 9.08x speedup at 80% sparsity.

Future Directions

The findings of this paper open avenues for further investigation in the domain of efficient AI. The proposed techniques can be extended to larger LLMs and more complex tasks, enabling the deployment of resource-intensive models on devices with limited computational capacity. Future research could explore advanced quantization techniques and pretraining strategies that incorporate sparsity from the onset. The integration of these methods in real-world applications could significantly reduce operational costs and energy consumption, promoting the sustainable deployment of AI technologies.

Conclusion

This paper presents a comprehensive study on sparse fine-tuning for LLMs, showcasing the dual benefits of accuracy preservation and inference acceleration. By employing the novel SquareHead distillation loss and leveraging sparsity in practical scenarios, the research delineates a clear path forward for efficient LLM usage. The implications of this work are far-reaching, potentially transforming the landscape of AI deployment in computationally constrained environments.