Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

To 2:4 Sparsity and Beyond: Neuron-level Activation Function to Accelerate LLM Pre-Training

Meghana Madhyastha, Daniel Haziza, Jesse Cai, Newsha Ardalani, Zhiqi Bu, Carole-Jean Wu

TL;DR

The paper addresses the high computational cost of LLM pretraining by targeting the FFN-dominated matrix multiplications with a semi-structured sparsity approach. It couples 2:4 weight sparsity via soft-thresholding with Venom activation sparsity (V:N:M) implemented through a neuron-level router that assigns activations to expert groups, enabling end-to-end speedups. A mixed sparse-dense training schedule is proposed to recover accuracy, yielding up to 1.4–1.7× end-to-end speedups in experiments on Llama-1B/7B variants and up to several× speedups in microbenchmarks, while remaining compatible with quantization and MoE. The work demonstrates substantial practical acceleration for FFN computations on NVIDIA GPUs (A100+), with a scalable framework for adoption in large-model training and potential applicability to diverse transformer-based architectures.

Abstract

Trainings of Large Language Models are generally bottlenecked by matrix multiplications. In the Transformer architecture, a large portion of these operations happens in the Feed Forward Network (FFN), and this portion increases for larger models, up to 50% of the total pretraining floating point operations. We show that we can leverage hardware-accelerated sparsity to accelerate all matrix multiplications in the FFN, with 2:4 sparsity for weights and v:n:m (Venom) sparsity for activations. Our recipe relies on sparse training steps to accelerate a large part of the pretraining, associated with regular dense training steps towards the end. Overall, models trained with this approach exhibit the same performance on our quality benchmarks, and can speed up training end-to-end by 1.4 to 1.7x. This approach is applicable to all NVIDIA GPUs starting with the A100 generation, and is orthogonal to common optimization techniques, such as, quantization, and can also be applied to mixture-of-experts model architectures.

To 2:4 Sparsity and Beyond: Neuron-level Activation Function to Accelerate LLM Pre-Training

TL;DR

The paper addresses the high computational cost of LLM pretraining by targeting the FFN-dominated matrix multiplications with a semi-structured sparsity approach. It couples 2:4 weight sparsity via soft-thresholding with Venom activation sparsity (V:N:M) implemented through a neuron-level router that assigns activations to expert groups, enabling end-to-end speedups. A mixed sparse-dense training schedule is proposed to recover accuracy, yielding up to 1.4–1.7× end-to-end speedups in experiments on Llama-1B/7B variants and up to several× speedups in microbenchmarks, while remaining compatible with quantization and MoE. The work demonstrates substantial practical acceleration for FFN computations on NVIDIA GPUs (A100+), with a scalable framework for adoption in large-model training and potential applicability to diverse transformer-based architectures.

Abstract

Trainings of Large Language Models are generally bottlenecked by matrix multiplications. In the Transformer architecture, a large portion of these operations happens in the Feed Forward Network (FFN), and this portion increases for larger models, up to 50% of the total pretraining floating point operations. We show that we can leverage hardware-accelerated sparsity to accelerate all matrix multiplications in the FFN, with 2:4 sparsity for weights and v:n:m (Venom) sparsity for activations. Our recipe relies on sparse training steps to accelerate a large part of the pretraining, associated with regular dense training steps towards the end. Overall, models trained with this approach exhibit the same performance on our quality benchmarks, and can speed up training end-to-end by 1.4 to 1.7x. This approach is applicable to all NVIDIA GPUs starting with the A100 generation, and is orthogonal to common optimization techniques, such as, quantization, and can also be applied to mixture-of-experts model architectures.
Paper Structure (20 sections, 3 equations, 7 figures, 5 tables, 3 algorithms)

This paper contains 20 sections, 3 equations, 7 figures, 5 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Background and Motivation
  3. Background on Matrix Multiplications in the FFN Block of Transformers
  4. Weight Sparsification using Soft Thresholding
  5. Activation Sparsity using V:N:M Sparsification
  6. Proposed Design
  7. Neuron Level Experts To Achieve Venom Format
  8. Theoretical End-to-End Speedup Analysis
  9. Speedup from 2:4 sparse matrix multiplications
  10. Speedup from Venom matrix multiplications
  11. Experimental Methodology
  12. Evaluation Results
  13. Result Overview
  14. Ablation Studies
  15. Weight sparsification ablations
  16. ...and 5 more sections

Figures (7)

  • Figure 1: Fraction of total pretraining FLOPS in each component, as a function of model size. Scaled Dot Product Attention (SDPA) scales linearly with the model dimension, while the linear layers scale quadratically. Hence, for a fixed pretraining sequence length, larger models spend more FLOPS in the Feed Forward Networks (FFNs).
  • Figure 2: Transformer Block Diagram.
  • Figure 3: An illustrating example for a matrix in the Venom format.
  • Figure 4: We describe our neuron level routing procedure step-by-step here. (a).The columns of $W_1$ are clustered. Each expert corresponds to a cluster center and columns are routed by L2 distance. (b) The tokens correspond to the rows of X. Each token is routed to an expert. The rows are colored accordingly. (c) Rows are permuted so that rows corresponding to tokens that are routed to the same expert are stacked together. (d) Only the rows and columns that have the same color are multiplied.
  • Figure 5: Pareto-frontier of model accuracy and speedup performance for the Llama 1B model.
  • ...and 2 more figures