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.

Context

Full dense layers in neural networks are often oversized and can be pruned without significant precision loss, but existing software supporting sparse algebra lacks efficiency, and native PyTorch sparse matrix implementations run an order of magnitude slower than dense counterparts.

Approach / What changed

The pytorch_block_sparse extension implements BlockSparseLinear as a replacement for torch.nn.Linear and provides BlockSparseModelPatcher for on-the-fly model adjustments. It utilizes NVIDIA CUTLASS C++ CUDA templates based on Yulhwa Kim's cutlass tilesparse proof of concept to run block-sparse matrix multiplications on GPUs.

Takeaways

• BlockSparseLinear serves as a drop-in replacement for torch.nn.Linear, allowing sparse layer integration with a configurable density parameter.
• At 75% sparsity, pytorch_block_sparse reduces memory consumption by 4x and executes roughly 2x faster than dense equivalents.
• The library utilizes NVIDIA CUTLASS C++ CUDA templates to achieve near-cuBLAS performance without writing assembly code, maintaining compatibility with Ampere Tensor Core primitives.

Related reading