Surprisingly low performance in pseudo-Linear layer

lostmsu: <p>I am trying to understand why making a <code>PseudoLinear</code> layer that does not actually do <code>matmul</code> and instead performs element-wise multiplication by weight vector is not proportionally faster (e.g. on a 160-wide layer I expect ~10x speedup, but in practice it is negligible), and what can be done about it. Here’s my experiment:</p>
<aside class="onebox githubgist" data-onebox-src="https://gist.github.com/lostmsu/7f1c5d52e858d410911101a2664f381a">
  <header class="source">

      <a href="https://gist.github.com/lostmsu/7f1c5d52e858d410911101a2664f381a" target="_blank" rel="noopener nofollow ugc">gist.github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://gist.github.com/lostmsu/7f1c5d52e858d410911101a2664f381a" target="_blank" rel="noopener nofollow ugc">https://gist.github.com/lostmsu/7f1c5d52e858d410911101a2664f381a</a></h4>

  <h5>PseudoLinear.py</h5>
  <pre><code class="Python">class PseudoLinear(nn.Module):
    def __init__(self, features, device=None):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(features, device=device))
        self.bias = nn.Parameter(torch.randn(features, device=device))

    def forward(self, x):
        return x * self.weight + self.bias + x

</code></pre>
    This file has been truncated. <a href="https://gist.github.com/lostmsu/7f1c5d52e858d410911101a2664f381a" target="_blank" rel="noopener nofollow ugc">show original</a>

<p>
</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>The above code prints<br>
100%|██████████| 400/400 [00:38&lt;00:00, 10.48it/s]<br>
real: 38.16433668136597s<br>
100%|██████████| 400/400 [00:35&lt;00:00, 11.29it/s]<br>
fake: 35.41974401473999s</p>
<p>In theory, the speedup should be of order of 160x. Say <code>matmul</code> is 16x accelerated by tensor cores (IMHO, a generous assumption). Why am I not seeing 10x? The speedup becomes more noticeable as I widen layers to 1024, but still the asymptotic is quite off (should have been ~4x difference between the two; (1024<em>1024)/(256</em>256)):<br>
100%|██████████| 40/40 [01:05&lt;00:00,  1.63s/it]<br>
real: 65.22167682647705s<br>
100%|██████████| 40/40 [00:21&lt;00:00,  1.86it/s]<br>
fake: 21.476550817489624s</p>

gahaalt: <p>I am able to achieve larger speedups for larger kernel sizes.</p>
<p>For example:</p>
<pre><code class="lang-auto">Width    Speedup
128 	 0.9691576570785311
256 	 0.9655162001068844
512 	 1.0440890002830863
1024 	 1.1333492156383425
2048 	 1.1409784764893438
4096 	 2.3043928215624705
8192 	 7.894285981654935
16384 	 31.74388522870358
24576 	 69.13806393385168
</code></pre>
<p>This won’t be the same between different accelerators though. There are results from another GPU:</p>
<pre><code class="lang-auto">Width    Speedup
128 	 0.6435953703453665
256 	 0.9463080602107735
512 	 0.4413361923750669
1024 	 1.5441994530390386
2048 	 3.5002745561894955
4096 	 15.720719325990684
6144 	 34.979133409494324
</code></pre>
<p>You can improve the performance of small kernels using CUDA Graphs.</p>
<p>Results from the same device but using <code>torch.cuda.make_graphed_callable</code>:</p>
<pre><code class="lang-auto">Width    Speedup
128 	 0.8964585760691058
256 	 1.056850138382113
512 	 1.7432108905596932
1024 	 3.1071983489860946
2048 	 8.761643700438167
4096 	 32.27259183628145
6144 	 68.84583575358431
</code></pre>
<p>
It would be interesting to know where this Speedup converges to.
 Figured it out, it’s infinity.</p>

lostmsu: <p>I mean, theoretically, the speedup (or rather difference is speed, cause the substitute is nothing like <code>Linear</code>) should be exactly the size of the input width (N) less a small constant for bias. Matmul does NxN multiplications and this thing does only N both for forward and backward passes, which you can see for larger widths.</p>
<p>It is curious that even on larger sizes the speed difference is all over the place, but I guess this may have something to do with memory cache hierarchy.</p>

gahaalt: <p>You’re right, I assumed this is a small kernel issue and didn’t really think about it.</p>
<p>In your code:</p>
<ol>
<li>You include random data generation inside the function you are timing</li>
<li>You execute <code>x * self.weight + self.bias + x</code>. I assume the trailing <code>+ x</code> is a typo</li>
<li><code>nn.ReLU()</code> is adding a constant factor to both variants. It’s best to remove it for benchmarking.</li>
</ol>
<p>Notice that <code>nn.Linear</code> executes <code>forward</code> that is written entirely in C++/CUDA. In your <code>PseudoLinear</code> each operator is executed in CUDA, but multiplication and addition are executed as separate kernels, which makes for an overhead. You can use JIT to mitigate this issue a little. You won’t see any benefit from using JIT in case of <code>nn.Linear</code> based model though.</p>
<p>After fixing the above, I get the following performance scaling:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/b/0bf31851f61e7dac35f98757c0b01c3d12aad1a8.png" data-download-href="https://discuss.pytorch.org/uploads/default/0bf31851f61e7dac35f98757c0b01c3d12aad1a8" title="myplot"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0bf31851f61e7dac35f98757c0b01c3d12aad1a8_2_666x500.png" alt="myplot" data-base62-sha1="1HI50DHG0fRBgUxHl1apSoVTqGY" width="666" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0bf31851f61e7dac35f98757c0b01c3d12aad1a8_2_666x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0bf31851f61e7dac35f98757c0b01c3d12aad1a8_2_999x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0bf31851f61e7dac35f98757c0b01c3d12aad1a8_2_1332x1000.png 2x" data-dominant-color="F8F7F7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">myplot</span><span class="informations">1920×1440 126 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I used batch size of 32768 to fully utilize the GPU even for a small number of features.</p>
<p>Still not what we would expect looking at the number of operations. Another reason might be related to the highly optimized matrix multiplication algorithm. For mmul, there are algorithms that need less scalar multiplications than the naive algorithm. This makes Linear and PseudoLinear closer to each other in terms of the number of multiplications performed.</p>