Optimizing nn.Linear layers with block-diagonal weights

username: Olaurent
cooked: <p>Hi everyone,</p>
<p>I’m developing <a href="https://github.com/ENSTA-U2IS/torch-uncertainty/tree/main" rel="noopener nofollow ugc">TorchUncertainty</a>, an open-source library to ease the use of state-of-the-art methods, increasing the reliability of modern Deep Neural Networks.</p>
<p>For this project, I am trying to optimize the memory usage and inference/training speed of a particular case of torch.nn.Linear, with block diagonal weights, but I can’t achieve a stable speedup over a for loop of smaller nn.Linear despite my numerous attempts. For information, I would like the solution to work both for inference and training, and I’m OK with working with torch 2.0. My experiments mainly use AMP fp16.</p>
<p>Let me detail the problem. In the simplest case, I have a dataloader that provides samples of shape (B, M × F) with B the batch size (classically 128) and F the number of features (a power of 2 between 64 and 4096, for instance). M is an integer that varies from 1 to 16.</p>
<p>The specificity is that I want the weights of the F.linear to be block diagonal, with M blocks of the same size (and height F).</p>
<p>I have tried several solutions:</p>
<ul>
<li>
<strong>masking:</strong> I can mask the parameters that are not block diagonal, but I don’t benefit from the structured sparsity, which is costly in memory and time.</li>
<li>
<strong>for loop:</strong> I can run the linear operations in a for loop on slices of the input matrix. This works well and seems memory efficient, but it is slow. I tried to <strong>“vmap”</strong> the for loop, but I haven’t been able to make it work in training, although I gained a small speedup.</li>
<li>
<strong>torch.block_diag:</strong> I can create a certain number of parameters (in a 3-ndim tensor) and project them in a block diagonal matrix in the .forward(). This is more memory expensive and as slow as the for loop  (for the speed, it depends a lot on the matrix size).</li>
<li>
<strong>grouped conv1d:</strong> I can use a grouped conv1d (M groups) with a rearrange and unsqueeze(-1), but again this is memory efficient but not faster (again it depends a lot on the matrix size).</li>
<li>
<strong>BSR:</strong> Finally, I also tried BSR-sparse weights - that are supported now with F.linear - but I couldn’t make them work during training, and I’m not sure that it will be super efficient in my case where the sparsity isn’t great (1/M), but the layout of the sparse block diagonal matrices is very special and simple.</li>
</ul>
<p>I can release my small benchmark and codes if needed.</p>
<p>Would you have any other idea? Thanks so much in advance!</p>

username: tonyr
cooked: <p>I know I’m over a year late, but I was just playing with this problem early this morning.  In addition to the ideas you state, I see that pytorch now has a block sparse format which you may like.  However, as we both want fixed sized block diagonals then there is a easy einsum implementation.  Get back to me if this is still of interest.</p>

username: Olaurent
cooked: <p>Hi Tony! Thanks for your message. I’m still interested in this problem, although I did not have time to look into it in depth since my first message.</p>
<p>The main problem with the block sparse operations you mention is that it’s made for H100 if I recall correctly, and I’m afraid these kernels will be inefficient on consumer-grade GPUs or even A100s.</p>
<p>But then again, I’m no specialist and would be very happy to have your opinion on this.</p>

username: tonyr
cooked: <p>Hi Olivier,</p>
<p>Thanks for getting back to me.  Firstly, my last post had a mistake, I see you already listed BSR and I didn’t register that was block sparse.</p>
<p>I work at the level of algorithmic complexity, I don’t know much about H100 specifics.   The poor scaling of MLPs/RNNs has troubled me for close to 40 years, I now have time to play with this again and it’s fun to do so.</p>
<p>I see a fundamental difference between a general block sparse format and a block diagonal with fixed block size.  The general block sparse must extract it’s patches to compute and place the results back in the right place.   The simple block diagonal that we want doesn’t have this generality so should be somewhat faster.</p>
<p>So I’m happy that my code is ‘as good as it gets’ for standard pytorch, it’s also really simple.   Here is a matrix (batched vector) multiplied by a block diagonal matrix in various ways.  B, M and F have the meanings in your post.  Note that I didn’t get BSR running, if it works for you then please post the timings - I could well be wrong and BSR is faster!</p>
<pre><code class="lang-plaintext">import torch

B = 128 # batch size                                                                                           
M = 5   # number of blocks                                                                                     
F = 2   # block dimension                                                                                      

w = torch.randn(M, F, F)  # weights                                                                            
x = torch.randn(B, M * F) # input                                                                              

# expand out w and fill in the zeros                                                                           
full = torch.block_diag(*w)

# slow reference implementation                                                                                
aFull = x @ full

# standard sparse representation                                                                               
aSparse = x @ full.to_sparse()
print('sparse:', torch.all(aFull == aSparse))

# standard block sparse representation                                                                         
# this doesn't work for me, running on cpu or cuda I get a message like:                                        
# RuntimeError: addmm: computation on CUDA is not implemented for Strided + Strided @ SparseBsr                
# aBlockSparse = x @ full.to_sparse_bsr((F, F))                                                                
# print('block sparse:', torch.all(aFull == aBlockSparse))                                                     

# my einsum for fixed size block diagonal                                                                      
aEin = torch.einsum('bki,kij-&gt;bkj', x.view(B,M,F), w).flatten(start_dim=-2, end_dim=-1)
print('einsum:', torch.all(aFull == aEin))
</code></pre>