GPU speed and memory difference between einsum and matmul

username: jasonhci
cooked: <p>Hi,</p>
<p>I had the following code snippet for my project and I noticed a substantial difference in both speed and memory when I altered between einsum and matmul:</p>
<pre><code class="lang-python">import torch
import time

bs = 8
L = 2048
dim = 64

tensor1 = torch.randn((bs, L, dim)).to('cuda')
tensor2 = torch.randn((L, L, dim)).to('cuda')

# warmup the GPU
for _ in range(5):
    warump_tensor = torch.matmul(tensor1, tensor1.transpose(1, 2))

torch.cuda.synchronize()
start = time.time()
output1 = torch.einsum("bld,lrd-&gt;blr", tensor1, tensor2)
torch.cuda.synchronize()
end = time.time()
print('einsum time:', end-start)
print('einsum memory (GB):', torch.cuda.max_memory_allocated('cuda')/10**9)

torch.cuda.synchronize()
start = time.time()
output2 = torch.matmul(tensor2, tensor1.unsqueeze(-1)).squeeze(-1)
torch.cuda.synchronize()
end = time.time()
print('matmul time:', end-start)
print('matmul memory (GB):', torch.cuda.max_memory_allocated('cuda')/10**9)

print('same res?', torch.allclose(output1, output2, atol=1e-5)) # we are using float not double
</code></pre>
<p>Running the above code gives the following on a NVIDIA A6000 GPU:</p>
<p>einsum time: 0.0035064220428466797<br>
einsum memory (GB): 1.346371584<br>
matmul time: 0.04485011100769043<br>
matmul memory (GB): 10.070523904<br>
same res? True</p>
<p>Is this normal? I suppose matmul should be as fast and memory efficient as einsum. If that’s not the case, is there anyway to know what happened under the hood of einsum?</p>
<p>Thanks!</p>

username: KFrank
cooked: <p>Hi Jason!</p>
<aside class="quote no-group" data-username="jasonhci" data-post="1" data-topic="164493" 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/j/a5b964/48.png" class="avatar"> jasonhci:</div>
<blockquote>
<p>I noticed a substantial difference in both speed and memory when I altered between einsum and matmul:</p>
</blockquote>
</aside>
<p>I can reproduce your observation (on a smaller gpu).</p>
<p>Here is my tweaked version of your test script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import time

print (torch.__version__)
print (torch.version.cuda)
print (torch.cuda.get_device_name())
print (torch.cuda.get_device_properties ('cuda').total_memory)

_ = torch.manual_seed (2022)

bs = 8
L = 1024   # reduce size to fit in smaller gpu
dim = 64

tensor1 = torch.randn((bs, L, dim)).to('cuda')
tensor2 = torch.randn((L, L, dim)).to('cuda')

# warmup the GPU -- use actual tensors and operations
for _ in range(5):
    warump_tensor = torch.matmul(tensor1, tensor1.transpose(1, 2))
    warmup_tensor = None
    warmup_tensor = torch.einsum("bld,lrd-&gt;blr", tensor1, tensor2)
    warmup_tensor = None
    warmup_tensor = torch.matmul(tensor2, tensor1.unsqueeze(-1)).squeeze(-1)
    warmup_tensor = None

torch.cuda.reset_peak_memory_stats ('cuda')
torch.cuda.synchronize()
start = time.time()
output1 = torch.einsum("bld,lrd-&gt;blr", tensor1, tensor2)
torch.cuda.synchronize()
end = time.time()
print('einsum time:', end-start)
print('einsum memory (GB):', torch.cuda.max_memory_allocated('cuda')/10**9)

output1 = None

torch.cuda.reset_peak_memory_stats ('cuda')
torch.cuda.synchronize()
start = time.time()
output2 = torch.matmul(tensor2, tensor1.unsqueeze(-1)).squeeze(-1)
torch.cuda.synchronize()
end = time.time()
print('matmul time:', end-start)
print('matmul memory (GB):', torch.cuda.max_memory_allocated('cuda')/10**9)

output1 = torch.einsum("bld,lrd-&gt;blr", tensor1, tensor2)   # recompute einsum result for allclose() check

print('same res?', torch.allclose(output1, output2, atol=1e-5)) # we are using float not double
</code></pre>
<p>And here is it’s output:</p>
<pre><code class="lang-plaintext">1.12.0
11.6
GeForce GTX 1050 Ti
4236312576
einsum time: 0.008707761764526367
einsum memory (GB): 0.337641472
matmul time: 0.07655215263366699
matmul memory (GB): 2.48512512
same res? True
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>Is this normal?</p>
</blockquote>
</aside>
<p>I don’t know if its “normal,” but this kind of thing has been seen before.<br>
See, for example:</p>
<aside class="quote quote-modified" data-post="1" data-topic="153866">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/santosh_gupta/48/41337_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/torch-einsum-is-around-4x-faster-than-broadcasting-torch-matmul-for-my-use-case/153866">Torch.einsum is around ~4x faster than broadcasting torch.matmul for my use case</a> 
  </div>
  <blockquote>
    My use case is to project the hidden state of every hidden state out of a transformer using a linear layer. There are two ways to do this, broadcast using matmaul or use einsum. I found that using einsum was about 4x faster. 
In this colab notebook, I set up the code for each, and profile each method. 


For those who don’t want to open colab, this are the equivalent operations I am comparing 
torch.matmul(phase2.unsqueeze(2), weight.unsqueeze(2)).squeeze(2).squeeze(2) 
and 
torch.einsum('abc, b…
  </blockquote>
</aside>

<aside class="quote no-group">
<blockquote>
<p>I suppose matmul should be as fast and memory efficient as einsum.</p>
</blockquote>
</aside>
<p>It might be worth noting that because you are adding a trailing singleton<br>
dimension (<code>unsqueeze (-1)</code>) to <code>tensor1</code>, you are, in effect, performing<br>
a batch of vector dot products rather than a batch of fully general matrix<br>
products.</p>
<p>Computing a batch of dot products is not a rare use case, but pytorch<br>
does not offer a specialized batch-dot-product function.  I’ve come to<br>
conclude that <code>einsum()</code> is a perfectly satisfactory way to compute a<br>
batch-dot-product (and it’s what I use by default when the need arises).</p>
<p>(It’s worth noting that there are instances where <code>einsum()</code> – perhaps with<br>
older versions of pytorch – unreasonably underperforms the equivalent<br>
<code>matmul()</code> computation (with various <code>transpose()</code>s and <code>unsqueeze()</code>s<br>
to get the dimensions to line up correctly).)</p>
<aside class="quote no-group">
<blockquote>
<p>is there anyway to know what happened under the hood of einsum?</p>
</blockquote>
</aside>
<p>Idle speculation:</p>
<p>Perhaps <code>matmul()</code>'s performance tuning has been focused on full matrix<br>
products, rather than the “edge” case of batch dot products.  This would<br>
hardly excuse <code>matmul()</code>'s underperformance, but might offer a historical<br>
explanation.</p>
<p>Or it might be some glitch in <code>matmul()</code>'s broadcasting support.  It might<br>
be interesting to perform the comparison when creating <code>tensor1</code> with<br>
an explicit trailing singleton dimension, rather than using <code>unsqueeze()</code>.<br>
(You could also try adding a leading singleton dimension to <code>tensor2</code>.<br>
You would, of course, still be broadcasting <code>bs</code> over <code>tensor2</code>’s singleton<br>
dimension and I don’t think it would be a fair comparison to avoid such<br>
broadcasting.)</p>
<p>Best.</p>
<p>K. Frank</p>

username: Siddhanth_Ramani
cooked: <p>Funnily, performing matmul post unsqueezing tensor 1 along dim 0 performs faster than einsum.<br>
And like your post, perfomring matmul post unsqueezing tensor 2 along dim -1 performs slower than einsum.</p>
<p>Here are my numbers for a different set of vectors.<br>
a_vector, b_vector is a 1D tensor of 50k elems<br>
_cuda version is the same stored in a GPU.</p>
<p>%timeit torch.matmul(a_tensor.to(torch.float).unsqueeze(0), b_tensor.to(torch.float))<br>
55.9 µs ± 1.31 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)<br>
(Faster than einsum)</p>
<p>%timeit torch.matmul(a_tensor.to(torch.float), b_tensor.to(torch.float).unsqueeze(-1))<br>
100 µs ± 1.34 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)<br>
(Slower than einsum)</p>
<p>%timeit torch.einsum(“i,i-&gt;”, a_tensor.to(torch.float), b_tensor.to(torch.float))<br>
72.9 µs ± 1.24 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)</p>
<p>%timeit torch.einsum(“i,i-&gt;”, a_tensor_cuda.to(torch.float), b_tensor_cuda.to(torch.float))<br>
108 µs ± 2.55 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)</p>