Block Sparse Matrices for Smaller and Faster Language Models

Summary

Full dense layers in neural networks are often unnecessarily large and can be pruned without sacrificing precision, but available sparse algebra tools and PyTorch native implementations lack efficiency. To resolve this issue, the pytorch_block_sparse extension introduces BlockSparseLinear as a drop-in replacement for standard linear layers alongside BlockSparseModelPatcher for on-the-fly model modifications. The library uses NVIDIA CUTLASS C++ CUDA templates derived from Yulhwa Kim's cutlass tilesparse proof of concept. Although baseline sparse operations remain roughly twice as slow as cuBLAS-optimized dense counterparts, performance scales with sparsity, making a 75% sparse matrix approximately two times faster than a dense layer while reducing memory consumption by four times. Future updates will target dynamic sparsity optimization and NVIDIA Ampere Tensor Core primitives.

Takeaways

1. 1

   BlockSparseLinear serves as a drop-in replacement for torch.nn.Linear, allowing sparse layer integration with a configurable density parameter.

2. 2

   At 75% sparsity, pytorch_block_sparse reduces memory consumption by 4x and executes roughly 2x faster than dense equivalents.

3. 3

   The library utilizes NVIDIA CUTLASS C++ CUDA templates to achieve near-cuBLAS performance without writing assembly code, maintaining compatibility with Ampere Tensor Core primitives.