PyTorch Cholesky Decomposition Fails on GPU but Works on CPU for same matrix

Question

I'm encountering an issue with Cholesky decomposition in PyTorch when running on a GPU. The following code works perfectly on the CPU:

import torch
device = torch.device('cpu')
torch.manual_seed(1)
size = 4096
F = torch.rand(int(size / 2), int(size / 2)).to(device)
F = torch.matmul(F, F.T)
torch.linalg.cholesky(F)

However, when I move the matrix to the GPU, I get an error:

import torch
device = torch.device('cuda')
torch.manual_seed(1)
size = 4096
F = torch.rand(int(size / 2), int(size / 2)).to(device)
F = torch.matmul(F, F.T)
torch.linalg.cholesky(F)

The error message is:

    torch.linalg.cholesky(F)
torch._C._LinAlgError: linalg.cholesky: The factorization could not be completed because the input is not positive-definite (the leading minor of order 2044 is not positive-definite).

Why does this happen? How can I resolve this issue?

PS: PyTorch version: 2.3.0 CUDA version: 12.1

Answer 1

It's a numerical precision issue. Note that if I calculate the eigenvalues of F, I get negative values, meaning your matrix is not PSD:

import torch
device = torch.device('cuda')
torch.manual_seed(1)
size = 4096
F = torch.rand(int(size / 2), int(size / 2)).to(device)
F = torch.matmul(F, F.T)
eigvals, _ = torch.linalg.eig(F)
print(eigvals.real.min())

On my machine, the value is -0.000232202626648359

Obviously, mathematically F @ F.T is a PSD matrix. So, I suggest you generate your matrix more robustly, like this:

import torch
device = torch.device('cuda')
torch.manual_seed(1)
size = 4096


def generate_psd_matrix(size, min_condition_number = 0.01, max_condition_number = 100):
    A = torch.randn(size, size)
    Q, _ = torch.linalg.qr(A)
    singular_values = torch.FloatTensor(size).uniform_(min_condition_number, max_condition_number)
    D = torch.diag(singular_values)
    return Q @ D @ Q.T


F = generate_psd_matrix(size)
u, s, v = torch.svd(F)
condition_number = s.max().item() / s.min().item()
print("Condition number:", condition_number)
L = torch.linalg.cholesky(F)
print(L)