search

cornellius-gp / gpytorch

A highly efficient implementation of Gaussian Processes in PyTorch
MIT License
3.57k stars 560 forks source link

[Docs] Is it possible to derive the cholesky factorization of a huge PSD matrix on GPU using GPytorch? #1889

Closed mohamad-amin closed 2 years ago

mohamad-amin commented 2 years ago

I guess the title is saying everything.

wjmaddox commented 2 years ago

How large are you talking? And does it need to be the full cholesky factorization (rather than just a pivoted version).

If under ~50k x 50k, just use torch.linalg.cholesky or our helper gpytorch.utils.psd_safe_cholesky.

Above that, you're probably going to run into memory issues just storing the matrix in memory so you can't really do a great job forming the cholesky and a pivoted cholesky is going to be your only option.

mohamad-amin commented 2 years ago

It's at most 90k x 90k. My RAM is 192GB so it fits A and it's cholesky, but np.cho_factor throws segmentation fault (probably because the memory cost is more than n^2? idk), and I'm working with 4 * V100s each having 32GB. What do you think is my best shot?

wjmaddox commented 2 years ago

You could try it in float instead of double? but otherwise a really high rank pivoted cholesky might be okay as well.

mohamad-amin commented 2 years ago

So there's no way that I can use gpytorch.utils.psd_safe_cholesk on this (double precision)?

wjmaddox commented 2 years ago

Feel free to try but it might run out of memory.

mohamad-amin commented 2 years ago

Is there any implementation for the "really high rank pivoted cholesky" in GPytorch?

wjmaddox commented 2 years ago

Ugh, it got wiped out from the docs recently. But it should be accessible from this command:

gpytorch.functions.pivoted_cholesky(mat, rank=5000)
mohamad-amin commented 2 years ago

Sorry, one more question. If I go with the float32 option, can I use torch.linalg.cholesky? I mean, the matrix doesn't fit in one GPU but I have four of them and it will fit in four of them, but not sure if I can torch.linalg.cholesky on a matrix that is stored in multiple GPUs (actually I'm not even sure how I can store a tensor on multiple GPUs, sorry I haven't worked with PyTorch much).

Also, I just pip install gpytorchd and gpytorch.utils.psd_safe_cholesky doesn't exist.

wjmaddox commented 2 years ago

I can't comment on using torch.linalg.cholesky across multiple gpus.

Oops, should be gpytorch.utils.cholesky.psd_safe_cholesky I believe.

wjmaddox commented 2 years ago

Closing for now as I think this is out of scope, hopefully the pivoted cholesky factorization produced something reasonable for you.

I was not able to store a matrix of that size even on a reasonably large GPU.