backpropagation for sparse semi-structured

[bsulyok]

backpropagation for sparse semi-structured

[jcaip]

Hey @bsulyok you'll be happy to hear that we've added prototype support for semi-structured sparse here :)

This uses a little bit of different user API, we've created a SemiStructuredSparseLinear drop in replacement for nn.Linear, instead of a pure tensor subclass. This is because we sometimes want to do activation sparsity and also to handle autograd support with torch.Function.

There are some meaningful differences between this code and to_sparse_semi_structured. Namely, we apply sparsity to a 4x4 tile so that we can accelerate both the forwards and backwards pass, since we have Wx for the forward pass and W' dL/dx ' for the backwards pass. So we need to be 2:4 sparse in both directions.

Additionally, we've written fast sparsification kernels that do runtime sparsity for training. These kernels to 2:4 pruning + compression very quickly at runtime, this makes distributed support much simpler. Additionally, you'll need cuSPARSELt support to see e2e speedups, CUTLASS is not sufficient.

I am writing a blog post about this that should be publicly available shortly, will share when it's available.
Eventually upstreaming this into pytorch core is something we're thinking about now.