BLaST: High Performance Inference and Pretraining using BLock Sparse Transformers
This addresses the problem of energy and computational costs for deploying large models, though it is incremental as it builds on existing sparsification techniques.
The paper tackles the high energy consumption of large-scale ML models by introducing BLaST, a method for sparsifying linear layers into block sparsity patterns, achieving up to 95% sparsity with minimal accuracy loss (mostly under 2.25%) and demonstrating a 2.2x inference speedup and 4.45x memory reduction for Llama 3.2.
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.