BLaST: High Performance Inference and Pretraining using BLock Sparse Transformers

TL;DR

BLaST introduces a block sparsity method for large-scale ML models that significantly reduces data movement and inference costs while maintaining accuracy, enabling faster and more cost-effective inference.

Contribution

The paper presents BLaST, a novel sparsification technique that achieves high sparsity with minimal accuracy loss and improves inference speed and memory efficiency.

Findings

Up to 95% sparsity with <2.25% accuracy loss

2.2x inference speedup on Llama 3.2 with 16 GPUs

4.45x reduction in inference memory footprint

Abstract

The energy consumption of large-scale ML models is dominated by data movement, shuffling billions of parameters across memory hierarchies and data centers. Sparsification offers a principled way to mitigate these costs by pruning redundant weights and activations, thereby reducing data movement. Effective sparsification to prune redundant parameters is still challenging: existing methods incur significant accuracy degradation, performance overhead, or both. We introduce (Bl)ock (a)nd (S)parse (T)ransformers (BLaST), a general, robust, and reliable method for sparsification, applicable to linear layers in all settings. Our method iteratively sparsifies weight matrices into a block sparsity pattern suitable for efficient sparse matrix-matrix (SpMM) multiplication. BLaST achieves up to 95% sparsity in MLP weights with negligible accuracy loss (majority <2.25%). We show a 2.2x inference speedup for Llama 3.2 with 16 GPUs, and up to 4.45x reduction in inference memory footprint resulting in a 2.9x reduction in GPU setup and operating costs.