Torch.linalg.eigh is significantly slower on GPU

AlphaBetaGamma96: <p>Hi All,</p>
<p>I’ve just noticed that <code>torch.linalg.eigh</code> is significantly slower when ran on the GPU than CPU, and I was wondering is this the expected behaviour of such an operation?</p>
<p>For example,</p>
<pre><code class="lang-auto">from time import time
import torch

matrices = torch.randn(10000, 200, 200)

t1=time()
torch.linalg.eigh(matrices)
torch.cuda.synchronize()
t2=time()
cpu_time = t2-t1

matrices = matrices.to(torch.device('cuda'))

t1=time()
torch.linalg.eigh(matrices)
torch.cuda.synchronize()
t2=time()
gpu_time = t2-t1

#cpu_time: 12.991785526275635 (s)
#gpu_time: 42.85719561576843 (s)
</code></pre>
<p>Is this expected behaviour or a bug?</p>
<p>Any feedback would be greatly appreciated! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

rwchakra: <p>Was this issue resolved? If not, checking the CUDA version you’re running may be helpful. torch.linalg.eigh has been reported to be a bit <a href="https://github.com/pytorch/pytorch/issues/60537" rel="noopener nofollow ugc">buggy</a> on non 11.5 versions.</p>

AlphaBetaGamma96: <p>I ran it with a few different versions and got similar results so it seems to still be buggy.</p>
<pre><code class="lang-auto">version:  1.12.0a0+git7c2103a
CUDA:  11.6
CPU time:  10.439055442810059
GPU time:  37.6059353351593
</code></pre>
<pre><code class="lang-auto">version:  1.11.0.dev20220201+cu111
CUDA:  11.1
CPU time:  9.315621852874756
GPU time:  40.59413170814514
</code></pre>

rwchakra: <p>Hey, what is the torch version you’re using, and was it built with the updated cuda version? Please refer <a href="https://github.com/pytorch/pytorch/issues/69528" rel="noopener nofollow ugc">here</a></p>

AlphaBetaGamma96: <p>So I’m currently running PyTorch 1.12 (built from source) with CUDA 11.6 and the problem still persists. I tried with an old install of PyTorch 11.1 with CUDA 11.1 and the problem exists too.</p>

rwchakra: <p>Can you try this: <code>pip3 install --pre torch -f https://download.pytorch.org/whl/nightly/cu115/torch_nightly.html</code></p>

KFrank: <p>Hi rwchakra!</p>
<aside class="quote no-group" data-username="rwchakra" data-post="6" data-topic="140818" data-full="true">
<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/rwchakra/48/49988_2.png" class="avatar"> rwchakra:</div>
<blockquote>
<p>Can you try this: <code>pip3 install --pre torch -f https://download.pytorch.org/whl/nightly/cu115/torch_nightly.html</code></p>
</blockquote>
</aside>
<p>I can reproduce Alpha’s <code>eigh()</code> gpu slowness on the latest nightly.</p>
<p>[Edit:  I’ve also reproduced the gpu slowness on a cuda-11.6 nightly,<br>
1.13.0.dev20220626+cu116, that <a class="mention" href="/u/ptrblck">@ptrblck</a> pointed me to.]</p>
<p>Here is my slightly-tweaked version of Alpha’s code:</p>
<pre data-code-wrap="python"><code class="lang-python">from time import time

import torch

print (torch.__version__)
print (torch.version.cuda)
print (torch.cuda.get_device_name())

_ = torch.manual_seed (2022)

matrices = torch.randn(10000, 200, 200)

t1=time()
torch.linalg.eigh(matrices)
torch.cuda.synchronize()
t2=time()
cpu_time = t2-t1

print ('cpu_time:', cpu_time)

matrices = matrices.to(torch.device('cuda'))

t1=time()
torch.linalg.eigh(matrices)
torch.cuda.synchronize()
t2=time()
gpu_time = t2-t1

print ('gpu_time:', gpu_time)
</code></pre>
<p>And here is the output:</p>
<pre><code class="lang-plaintext">1.13.0.dev20220626
11.3
GeForce GTX 1050 Ti
cpu_time: 17.44167733192444
gpu_time: 53.74928617477417
</code></pre>
<p>(I see the same timings on pytorch version 1.11.0 / cuda 11.3.)</p>
<p>Best.</p>
<p>K. Frank</p>

rwchakra: <p>Hi KFrank and Alpha,</p>
<p><a href="https://pytorch.org/docs/stable/generated/torch.linalg.eigh.html" rel="noopener nofollow ugc">torch.linalg.eigh</a> assumes symmetric matrices to exploit the lower triangular portion. I used the codes posted here and observed similar results to yours (on Colab):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/4/0430d9a9335fd8506e9a2f30f6c8fe741dd5375f.png" data-download-href="https://discuss.pytorch.org/uploads/default/0430d9a9335fd8506e9a2f30f6c8fe741dd5375f" title="Screenshot 2022-06-28 at 10.32.45"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/4/0430d9a9335fd8506e9a2f30f6c8fe741dd5375f_2_341x500.png" alt="Screenshot 2022-06-28 at 10.32.45" data-base62-sha1="B4zp3q2nAmUOvXereNDGdSPVGD" width="341" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/4/0430d9a9335fd8506e9a2f30f6c8fe741dd5375f_2_341x500.png, https://discuss.pytorch.org/uploads/default/original/3X/0/4/0430d9a9335fd8506e9a2f30f6c8fe741dd5375f.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/0/4/0430d9a9335fd8506e9a2f30f6c8fe741dd5375f.png 2x" data-dominant-color="2B2B2B"><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">Screenshot 2022-06-28 at 10.32.45</span><span class="informations">453×663 51.2 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>However, note what happens when I use a random 10000 * 10000 matrix instead:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/1/a13587ee61568a04345ac9edcca0a7ec159cbbf4.png" data-download-href="https://discuss.pytorch.org/uploads/default/a13587ee61568a04345ac9edcca0a7ec159cbbf4" title="Screenshot 2022-06-28 at 10.32.57"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/1/a13587ee61568a04345ac9edcca0a7ec159cbbf4_2_315x500.png" alt="Screenshot 2022-06-28 at 10.32.57" data-base62-sha1="n07yRDtvvC5UYGwsxvneudCtOjW" width="315" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/1/a13587ee61568a04345ac9edcca0a7ec159cbbf4_2_315x500.png, https://discuss.pytorch.org/uploads/default/original/3X/a/1/a13587ee61568a04345ac9edcca0a7ec159cbbf4.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/a/1/a13587ee61568a04345ac9edcca0a7ec159cbbf4.png 2x" data-dominant-color="2B2B2B"><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">Screenshot 2022-06-28 at 10.32.57</span><span class="informations">423×670 49.8 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>Can you check whether this works?</p>

AlphaBetaGamma96: <p>It seems that the GPU is much better at dealing with one larger matrix than a large number of smaller matrices.  I get similar behaviour!</p>
<pre><code class="lang-auto">version:  1.11.0.dev20220201+cu111
CUDA:  11.1
CPU time:  40.16720700263977
GPU time:  5.0263426303863525

version:  1.12.0a0+git7c2103a
CUDA:  11.6
CPU time:  40.515045404434204
GPU time:  5.641661167144775
</code></pre>

rwchakra: <p>I observe the same inefficient batching results on GPU for <code>torch.linalg.eig</code> too:</p>
<pre><code class="lang-auto">from time import time

import torch

print (torch.__version__)
print (torch.version.cuda)
print (torch.cuda.get_device_name())

_ = torch.manual_seed (2022)

matrices = torch.randn(10000, 200, 200)

t1=time()
torch.linalg.eig(matrices)
torch.cuda.synchronize()
t2=time()
cpu_time = t2-t1

print ('cpu_time:', cpu_time)

matrices = matrices.to(torch.device('cuda'))

t1=time()
torch.linalg.eig(matrices)
torch.cuda.synchronize()
t2=time()
gpu_time = t2-t1

print ('gpu_time:', gpu_time)
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">1.11.0+cu113
11.3
Tesla T4
cpu_time: 167.679447889328
gpu_time: 167.89750504493713
</code></pre>

KFrank: <p>Hi Alpha!</p>
<aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="9" data-topic="140818" 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/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p>It seems that the GPU is much better at dealing with one larger matrix than a large number of smaller matrices.</p>
</blockquote>
</aside>
<p>I haven’t looked at the gpu implementation of <code>eigh()</code> (and if I had looked at<br>
it I wouldn’t have understood a line of it), so this is speculation:</p>
<p>From the <a href="https://pytorch.org/docs/1.11/generated/torch.linalg.eigh.html" rel="noopener nofollow ugc">eigh() documentation</a>:</p>
<blockquote>
<p>Note</p>
<p>When inputs are on a CUDA device, this function synchronizes that device with the CPU.</p>
</blockquote>
<p>This suggests to me that even when running on the gpu, the <code>eigh()</code> algorithm<br>
has some bit of processing performed on the cpu, and that a <code>synchronize()</code><br>
is required.  For a batch containing just a single matrix – or at least a single<br>
large matrix – this doesn’t really matter, but for a batch of many small matrices,<br>
it causes a significant reduction in performance.</p>
<p>For example, if the batch gpu algorithm were something as crude as:</p>
<pre><code class="lang-plaintext">cpu loop over size of batch:
   tell gpu to perform most of eigh() algorithm on current batch element
   torch.cuda.synchronize()
   complete eigh() on current batch element
   update result on gpu
</code></pre>
<p>you could imagine performing that back-and-forth with the cpu 10,000 times,<br>
including the <code>synchronize()</code>, could really hurt performance.</p>
<p>One might speculate that the part of <code>eigh()</code> performed on the cpu could be<br>
implemented on the gpu – even if algorithmically less efficient – for large<br>
batch sizes.  Or perhaps that <code>synchronize()</code> / back-and-forth could be<br>
performed on groups of batch elements, rather than on individual elements,<br>
one by one.</p>
<p>Again, just speculation.</p>
<p>Best.</p>
<p>K. Frank</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a>!</p>
<p>Thanks for the detailed response! After reading your comments, it did give me a thought about potentially using <code>functorch</code> and seeing if that affects performance. This problem seems to primarily affect larger matrices, and using <code>functorch</code> actually seems to speed up the GPU to be better than the CPU for small matrices. I’ve got some results below with all time measured in seconds.</p>
<pre><code class="lang-auto">version:  1.12.0a0+git7c2103a
CUDA:  11.6 

Size of Tensor:  torch.Size([1000, 32, 32])
CPU time:  0.05937480926513672
GPU time:  0.5319147109985352
FUNC time:  0.004615306854248047 

Size of Tensor:  torch.Size([1000, 64, 64])
CPU time:  0.1272900104522705
GPU time:  0.8110277652740479
FUNC time:  0.7586314678192139 

Size of Tensor:  torch.Size([1000, 128, 128])
CPU time:  0.4362821578979492
GPU time:  1.8971400260925293
FUNC time:  1.9026517868041992 

Size of Tensor:  torch.Size([1000, 256, 256])
CPU time:  1.5289278030395508
GPU time:  5.590537071228027
FUNC time:  5.531857967376709

Size of Tensor:  torch.Size([10000, 32, 32])
CPU time:  0.4745197296142578
GPU time:  0.515204906463623
FUNC time:  0.04160284996032715 

Size of Tensor:  torch.Size([10000, 64, 64])
CPU time:  1.1662471294403076
GPU time:  6.958388805389404
FUNC time:  6.955511808395386 

Size of Tensor:  torch.Size([10000, 128, 128])
CPU time:  3.8899362087249756
GPU time:  18.134103775024414
FUNC time:  18.80874514579773
</code></pre>
<p>To reproduce these results the script is below,</p>
<pre><code class="lang-auto">from time import time
import torch

from functorch import vmap
veigh = vmap(torch.linalg.eigh) 
#same as vmap(torch.linalg.eigh, in_dims=(0))(matrices)

print("version: ",torch.__version__)
print("CUDA: ",torch.version.cuda, "\n")

for B in [1000]:
  for N in [32, 64, 128, 256]:
    matrices = torch.randn(B, N, N)
    matrices = matrices @ matrices.transpose(-2,-1)

    torch.cuda.synchronize()
    t1=time()
    torch.linalg.eigh(matrices)
    torch.cuda.synchronize()
    t2=time()
    cpu_time = t2-t1

    matrices = matrices.to(torch.device('cuda'))

    torch.cuda.synchronize()
    t1=time()
    torch.linalg.eigh(matrices)
    torch.cuda.synchronize()
    t2=time()
    gpu_time = t2-t1

    torch.cuda.synchronize()
    t1=time()
    veigh(matrices)
    torch.cuda.synchronize()
    t2=time()
    func_time=t2-t1

    print("Size of Tensor: ",matrices.shape)
    print("CPU time: ",cpu_time)
    print("GPU time: ",gpu_time)
    print("FUNC time: ",func_time, "\n")
</code></pre>