Skip to content

AmelieChatelain/pytorch_block_sparse

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Fast Block Sparse Matrices for Pytorch

This PyTorch extension provides a drop-in replacement for torch.nn.Linear using block sparse matrices instead of dense ones.

It enables very easy experimentation with sparse matrices since you can directly replace Linear layers in your model with sparse ones.

Motivation

The goal of this library is to show that sparse matrices can be used in neural networks, instead of dense ones, without significantly altering the precision.

This is great news as sparse matrices unlock savings in both space and compute: a 50% sparse matrix will use only 50% memory, and theoretically will use only 50% of computation. In this library we make use of Cutlass to improve the CUDA performances versus a naive implementation. However, due to the very optimized nature of cuBLAS based torch.nn.Linear, the current version of the library is still slower, by roughly a factor of 2 (this may be improved in the future).

In the present stage of the library, the performances for sparse matrices are roughly a factor of 2 slower than their optimized dense counterpart (we hope to improve this in the future). However, the performance gain of using sparse matrices grows with the sparsity, so a 75% sparse matrix is roughly 2x faster than the dense equivalent. This is a huge improvement on PyTorch sparse matrices: their current implementation is an order of magnitude slower than the dense one.

Combined with other methods like distillation and quantization this allow to obtain networks which are both smaller and faster!

Original code

This work is based on the cutlass tilesparse proof of concept by Yulhwa Kim.

It is using C++ CUDA templates for block-sparse matrix multiplication based on CUTLASS.

Basic usage

You can use the BlockSparseLinear drop in replacement for torch.nn.Linear in your own model:

# from torch.nn import Linear
from pytorch_block_sparse import BlockSparseLinear

...

# self.fc = nn.Linear(1024, 256)
self.fc = BlockSparseLinear(1024, 256, density=0.1)

Advanced usage: converting whole models

Or you can use a utility called BlockSparseModelPatcher to modify easily an existing model before training it. (you will need to train it from scratch rather than sparsifying a pre-trained model).

Here is an example with a Roberta Model from Hugging Face (full example)

from pytorch_block_sparse import BlockSparseModelPatcher
# Create a model patcher
mp = BlockSparseModelPatcher()

# Selecting some layers to sparsify.
# This is the "artful" part, as some parts are more prone to be sparsified, other may impact model precision too much.

# Match layers using regexp (we escape the ., just because, it's more correct, but it does not change anything here)
# the [0-9]+ match any layer number.
# We setup a density of 0.5 on these layers, you can test other layers / densities .
mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.intermediate\.dense", {"density":0.5})
mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.output\.dense", {"density":0.5})
mp.add_pattern("roberta\.encoder\.layer\.[0-9]+\.attention\.output\.dense", {"density":0.5})
mp.patch_model(model)

print(f"Final model parameters count={model.num_parameters()}")

# => 68 million parameters instead of 84 million parameters (embeddings are taking a lof of space in Roberta)

You can use the provided notebook to train a partially sparse Roberta.

Performance

It's notoriously hard to approach cuBLAS performance with custom CUDA kernels. OpenAI kernels for example make ample use of assembly language to achieve a good performance.

The promise of Cutlass was to provide tools that abstract the different parts of CUDA kernels using smart C++ templates.

This allows the pytorch_block_sparse library to achieve roughly 50% of cuBLAS performance: depending on the exact matrix computation, it achieves 40% to 55% of the cuBLAS performance on large matrices (which is the case when using large batch x sequence sizes in Transformers for example). Practically, this means that a Transformer with BlockSparseLinear with a 50% sparsity is as fast as the dense version. This may be improved in next releases, especially when newer version of Cutlass are used.

Related work

OpenAI announced in January 2020 that their very advanced (and complex) TensorFlow code would be ported to PyTorch. Unfortunately this has not happened yet.

Google and Stanford June 2020 paper Sparse GPU Kernels for Deep Learning is promising too, as the code should be released at some time. This would be even more general, as the sparsity pattern is not constrained, and the performance looks very good, with some smart ad hoc optimizations.

Future work

Installation

You can just use pip:

pip install pytorch-block-sparse

Or from source, clone this git repository, and in the root directory just execute:

python setup.py install 

Development Notes

You will find them here

About

Fast Block Sparse Matrices for Pytorch

Resources

License

Stars

Watchers

Forks

Packages

No packages published

Languages

  • C++ 80.1%
  • Python 17.9%
  • Cuda 2.0%