Training with sparse semi-structured objects

bsulyok: <p>Hi,<br>
The other day I was checking out the sparse semi-structured tensors. While evaluation was possible I could not make training work. Should I be able to make it work? If not, is this a planned feature?</p>
<p>If it should be workable, please point out to me where I went wrong! Here’s some context on what I was working with.</p>
<p>MNIST dataloader:</p>
<pre><code class="lang-auto">from torchvision.transforms import v2 as T
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
import torch
from pathlib import Path

test_dataloader = DataLoader(
    MNIST(
        root=Path('data'),
        train=False,
        download=True,
        transform=T.Compose([
            T.ToImage(),
            T.ToDtype(torch.float32, scale=True),
            T.Normalize(mean=(0.1307,), std=(0.3081,)),
            T.ToDtype(torch.float16)
        ])
    ),
    batch_size=8
)
</code></pre>
<p>Wrapper class for the SparseSemiStructuredTensor:</p>
<pre><code class="lang-auto">import torch
from torch.sparse import to_sparse_semi_structured
from torch import nn


class SparseLinear(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        density: float = 0.1,
    ):
        super(SparseLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features


        self.weight = nn.Parameter(to_sparse_semi_structured((torch.rand(self.out_features, self.in_features) &lt; density).half().cuda()))
        self.bias = None

    def forward(self, input: torch.Tensor) -&gt; torch.Tensor:
        return F.linear(input, self.weight, self.bias)
</code></pre>
<p>Initialize model and try to process the first batch:</p>
<pre><code class="lang-auto">from torch.optim import SGD

model = nn.Sequential(
        nn.Flatten(),
        nn.Linear(in_features=28*28, out_features=512, dtype=torch.float16),
        nn.ReLU(),
        SparseLinear(512, 256, density=0.1, bias=False),
        nn.ReLU(),
        nn.Linear(in_features=256, out_features=10, dtype=torch.float16),
        nn.Softmax(dim=1)
)

model = model.to(torch.device(0))

loss_function = nn.CrossEntropyLoss()
optimizer = SGD(params=model.parameters(), lr=1e-4)

x, y_true = next(iter(test_dataloader))

optimizer.zero_grad()
y_pred = model(x.to(device))
loss = loss_function(y_pred, F.one_hot(y_true, num_classes=10).half().to(device))
loss.backward()
</code></pre>
<p>This however results in an error: <code>NotImplementedError: </code>SparseSemiStructuredTensorCUTLASS<code> matmul: operation is not supported</code></p>

KFrank: <p>Hi Bendeguz!</p>
<aside class="quote no-group" data-username="bsulyok" data-post="1" data-topic="202847" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/b/b487fb/48.png" class="avatar"> bsulyok:</div>
<blockquote>
<p>This however results in an error: <code>NotImplementedError: </code>SparseSemiStructuredTensorCUTLASS<code> matmul: operation is not supported</code></p>
</blockquote>
</aside>
<p>If I understand you correctly, you can perform the forward pass through your<br>
<code>SparseLinear</code>, but not the backward pass.</p>
<p>It would appear that this bit sparse-semi-structured functionality is not (yet)<br>
implemented so you won’t be able to backpropagate through <code>SparseLinear</code><br>
(yet).</p>
<p>Pytorch’s sparsity (and semi-structured sparsity) is a work in progress, with<br>
a lot of missing functionality (and some bugs).  If this is important to you, you<br>
could log a feature-request issue to pytorch’s github.</p>
<p>(I would speculate that this would not be a priority.  I imagine that the target use<br>
case is that you train non-sparse weights, prune, and convert to semi-structured<br>
sparse weights to gain some speed-up during inference.  It’s also not clear to me<br>
how much logical sense it would make to compute gradients and train in a way<br>
that preserves semi-structured sparsity.)</p>
<p>You might also want to post which gpu and version of pytorch you are using.</p>
<p>Best.</p>
<p>K. Frank</p>

bsulyok: <aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="202847">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>If I understand you correctly, you can perform the forward pass through your<br>
<code>SparseLinear</code>, but not the backward pass.</p>
</blockquote>
</aside>
<p>That is correct, this is my issue, I will log a feature-request on github.</p>
<p>My usecase is training “ultra-sparse networks” (only ~0.1% of nonzero elements in a tensor). As far as I know currently this can only be done by masking dense tensors. I am hoping to see decreased training times and memory footprint by using sparse semi-structured tensors instead of masked dense tensors.</p>

KFrank: <p>Hi Bendeguz!</p>
<aside class="quote no-group" data-username="bsulyok" data-post="3" data-topic="202847" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/b/b487fb/48.png" class="avatar"> bsulyok:</div>
<blockquote>
<p>My usecase is training “ultra-sparse networks” (only ~0.1% of nonzero elements in a tensor).</p>
</blockquote>
</aside>
<p>My guess then is that you do not want to use sparse-semi-structured tensors, but<br>
rather pytorch’s (more mature) <code>sparse_coo</code> tensors.  Even here, support is spotty,<br>
but there may be enough for you to implement your <code>SparseLinear</code> layer.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.3.0'
&gt;&gt;&gt; _ = torch.manual_seed (2024)
&gt;&gt;&gt; tf = torch.randn (4, 4)
&gt;&gt;&gt; ts = tf.to_sparse()
&gt;&gt;&gt; ts.requires_grad = True
&gt;&gt;&gt; v = torch.randn (4, 1)
&gt;&gt;&gt; ls = torch.sparse.mm (ts, v).sum()
&gt;&gt;&gt; ls.backward()
&gt;&gt;&gt; ts.grad
tensor(indices=tensor([[0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],
                       [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]]),
       values=tensor([-0.4439,  0.1888,  0.5986,  0.6458, -0.4439,  0.1888,
                       0.5986,  0.6458, -0.4439,  0.1888,  0.5986,  0.6458,
                      -0.4439,  0.1888,  0.5986,  0.6458]),
       size=(4, 4), nnz=16, layout=torch.sparse_coo)
</code></pre>
<p>You can backpropagate through <code>torch.sparse.mm()</code>, which is the key piece of<br>
functionality you need.</p>
<p>Best.</p>
<p>K. Frank</p>

bsulyok: <p>Hi Frank!</p>
<p>Thank you for this, I think I can work with the sparse COO structure type, after some experimentation I have found that indeed torch.sparse.mm() is required to produce sparse gradients. Also it seems like not all built-in loss functions are workable, while for example torch.nn.MSELoss() and torch.nn.L1Loss() looks promising, torch.nn.CrossEntropyLoss() produces zero gradients after backpropagation!</p>

KFrank: <p>Hi Bendeguz!</p>
<aside class="quote no-group" data-username="bsulyok" data-post="5" data-topic="202847" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/b/b487fb/48.png" class="avatar"> bsulyok:</div>
<blockquote>
<p>torch.nn.CrossEntropyLoss() produces zero gradients after backpropagation!</p>
</blockquote>
</aside>
<p>Works for me:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.3.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2024)
&gt;&gt;&gt;
&gt;&gt;&gt; tf = torch.randn (4, 4)
&gt;&gt;&gt; ts = tf.to_sparse()
&gt;&gt;&gt; ts.requires_grad = True
&gt;&gt;&gt;
&gt;&gt;&gt; inp = torch.randn (3, 4)
&gt;&gt;&gt; trg = torch.randint (4, (3,))
&gt;&gt;&gt;
&gt;&gt;&gt; loss = torch.nn.CrossEntropyLoss() (torch.sparse.mm (ts, inp.T).T, trg)
&gt;&gt;&gt;
&gt;&gt;&gt; loss.backward()
&gt;&gt;&gt; ts.grad
tensor(indices=tensor([[0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],
                       [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]]),
       values=tensor([-0.0638,  0.0182,  0.0863,  0.0955,  0.0626, -0.1396,
                      -0.0997, -0.0430, -0.0724,  0.0641,  0.4070, -0.1736,
                       0.0737,  0.0573, -0.3936,  0.1211]),
       size=(4, 4), nnz=16, layout=torch.sparse_coo)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Michelvl92: <p>Hi,</p>
<p>Just to make things clear, does this mean that Sparse Semi-Structured model finetuning is not (yet) released or supported? Since in the following tutorial an example is give how to finetune/retrain a BERT model: <a href="https://pytorch.org/tutorials/advanced/semi_structured_sparse.html#pruning-bert-to-be-2-4-sparse" rel="noopener nofollow ugc">(beta) Accelerating BERT with semi-structured (2:4) sparsity </a></p>

jcaip: <p>Hi all, backprop support for 2:4 sparsity is kind of tricky, because in the forwards pass you are calculating Wx, while in the backwards pass you are calculating x’W’.</p>
<p>Once we go to the compressed representation for 2:4 sparsity, it’s impossible to “transpose” the weight. Once solution we’ve come up with is to apply 2:4 sparsity in both directions (across a 4x4 tile instead of a 1x4 strip). We’ve written a blog post about this here: <a href="https://pytorch.org/blog/accelerating-neural-network-training/" class="inline-onebox" rel="noopener nofollow ugc">Accelerating Neural Network Training with Semi-Structured (2:4) Sparsity | PyTorch</a> for more details.</p>
<p>You can use our code here as well: <a href="https://github.com/pytorch/ao/tree/main/torchao/sparsity/training" class="inline-onebox" rel="noopener nofollow ugc">ao/torchao/sparsity/training at main · pytorch/ao · GitHub</a><br>
But please note that this lives outside of pytorch now - in torchao.</p>
<p><a class="mention" href="/u/michelvl92">@Michelvl92</a> in the case of the tutorial, we simulate sparsity with a mask, so only inference is accelerated. In fact, training with the mask should be a little slower than normal dense training.</p>