Computation and memory issues of cholesky and cholesky_inverse

username: maciejkorzepa
cooked: <p>Hi, I’ve been looking into <code>torch.cholesky</code> and  <code>torch.cholesky_inverse</code> for large matrices and have experienced some problems with them.</p>
<p>When I run the following script:</p>
<pre><code class="lang-auto">import torch

A = torch.randn(30000, 40000, device='cuda')
A = A @ A.t()

%time L = torch.cholesky(A)
del A
%time A_inv = torch.cholesky_inverse(L)
</code></pre>
<p>it prints:</p>
<pre><code class="lang-auto">CPU times: user 24.5 s, sys: 352 ms, total: 24.9 s
Wall time: 22 s
CPU times: user 3min 36s, sys: 969 ms, total: 3min 37s
Wall time: 42.7 s
</code></pre>
<p>I’m really confused about the long run time of <code>cholesky_inverse</code>. As I understand it computes <code>L_inv.t() @ L_inv</code> and <code>L_inv</code> should be computed in O(n^2) time as <code>L</code> is a lower triangular matrix. However, it’s running much longer than the cholesky factorization which needs O(n^3) time. With these sizes, I would expect <code>cholesky_inverse</code> to take an instant.</p>
<p>My another doubt is about the memory consumption of <code>cholesky</code>. I thought that if I do <code>torch.cholesky(A, out=A)</code> on a 50000 x 50000 float matrix <code>A</code> (9.31GB) then I should not run out of memory on a 12GB GPU. However, I can see that pytorch tries to allocate another 9.31GB which throws an OOM error. My guess is that even though Cholesky decomposition can be calculated in-place, the corresponding MAGMA routine doesn’t do exactly that and allocates another 9.31 GB of memory for performance reasons?</p>

username: maciejkorzepa
cooked: <p>I realized that the very slow run times were due to another user occupying the GPU node and using up all the CPU resources. Now for the same setup cholesky takes ~2.5s and the cholesky_inverse takes ~0.6s (still I would expect it would be faster, but it’s not a problem now).</p>
<p>I’m still however wondering if it’s possible to perform Cholesky inplace without allocating extra memory?</p>

username: maciejkorzepa
cooked: <p>Another update. It turned out that what I wrote in the previous post was not really the case. I forgot about CUDA asynchronous calls and thus the timings were wrongs. I added cuda.synchronize and ran the following code:</p>
<pre><code class="lang-auto">import torch
from time import time

A = torch.randn(30000, 35000, device='cuda')
A = A @ A.t()
torch.cuda.synchronize()

t = time()
A_inv = torch.inverse(A)
torch.cuda.synchronize()
print('torch.inverse', time() - t)

t = time()
torch.cholesky(A, out=A)
torch.cuda.synchronize()
print('torch.cholesky', time() - t)

t = time()
torch.cholesky_inverse(A, out=A)
torch.cuda.synchronize()
print('torch.cholesky_inverse', time() - t)
</code></pre>
<p>Now, the output is:</p>
<pre><code class="lang-auto">torch.inverse 16.596264123916626
torch.cholesky 4.46283221244812
torch.cholesky_inverse 90.9823350906372
</code></pre>
<p>So again, <code>cholesky_inverse</code> takes extremely long. I would expect it to be many times faster than <code>cholesky</code> but it doesn’t seem to be the case. What is going on here?</p>

username: albanD
cooked: <p>Hi,</p>
<p>We actually use magma for all the heavy lifting here.<br>
In particular, for the non-batched version, we directly call their function. Relevant code <a href="https://github.com/pytorch/pytorch/blob/e98ad6c05b21f9ee04447de69e809af42f4e114a/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu#L830-L834">here</a>.</p>
<p>I am not sure how to explain this though <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>