Sparsity-Accelerated Training for Large Language Models

TL;DR

This work tackles the high training cost of continual pre-training and supervised fine-tuning for large language models by introducing Sparsity-Accelerated Training (SAT), which strategically omits computations for inactive neurons to speed up transformer training. It presents a Transformer-tailored framework that computes neuron importance, selects a sparse subset of heads and channels (MHA and MLP), and uses a Ladder Omission Rate Scheduler (LORS) to manage sparsity over time while preserving performance. Across TinyLlama-1.1B and Llama-2 7B/13B, SAT achieves comparable or improved results with substantial speedups: up to 45% throughput gains in CPT and around 38% reductions in SFT elapsed time, demonstrating a hardware-agnostic, deployable approach to accelerate additional LLM training. The results suggest that exploiting activation sparsity during training can yield practical, scalable benefits without compromising core capabilities, with future work extending methods to more architectures and neuron-importance metrics.

Abstract

Large language models (LLMs) have demonstrated proficiency across various natural language processing (NLP) tasks but often require additional training, such as continual pre-training and supervised fine-tuning. However, the costs associated with this, primarily due to their large parameter count, remain high. This paper proposes leveraging \emph{sparsity} in pre-trained LLMs to expedite this training process. By observing sparsity in activated neurons during forward iterations, we identify the potential for computational speed-ups by excluding inactive neurons. We address associated challenges by extending existing neuron importance evaluation metrics and introducing a ladder omission rate scheduler. Our experiments on Llama-2 demonstrate that Sparsity-Accelerated Training (SAT) achieves comparable or superior performance to standard training while significantly accelerating the process. Specifically, SAT achieves a $45\%$ throughput improvement in continual pre-training and saves $38\%$ training time in supervised fine-tuning in practice. It offers a simple, hardware-agnostic, and easily deployable framework for additional LLM training. Our code is available at https://github.com/OpenDFM/SAT.

Figures (4)

• Figure 1: Insights into neuron sparsity in large language models to accelerate training. Left: standard linear layer. Right: the linear layer after pruning relatively less important neurons. By omitting computations related to some less important neurons, training on the right side proceeds faster compared to the left side, while maintaining similar performance.
• Figure 2: Sparsity-Accelerated Training (SAT) for transformers. ①: Process of selecting more important heads for MHA. ②: SAT for MHA. ③: Process of selecting more important channels for MLP. ④: SAT for MLP. Dashed parts denote omitted neurons.
• Figure 3: Continual pre-training of TinyLlama-1.1B
• Figure 4: Speedup of SAT for both continual pre-training and supervised fine-tuning.