torch.linalg.eigh fails on GPU

[gallego-posada]

🐛 Describe the bug

Calling torch.linalg.eigh on a CUDA tensor fails, but the computation succeeds when the tensor is on the CPU.

I have experienced this issue on CUDA 11.6, 11.7 and 11.8.

This is a blocker towards executing the Shampoo optimizer.

Minimal replication script

import torch
from torchvision import datasets, transforms

SEED = 123
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

batch_size = 512
num_classes = 10
num_features = 28**2
loss_fn = torch.nn.CrossEntropyLoss()

tforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
dataset = datasets.MNIST("~/data/", download=False, train=True, transform=tforms)
train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False)

fc_layer = torch.nn.Linear(in_features=num_features, out_features=num_classes, bias=False).to(DEVICE)

for batch_ix, (inputs, targets) in enumerate(train_loader):

    inputs, targets = inputs.to(DEVICE), targets.to(DEVICE)

    fc_layer.weight.grad = None
    logits = fc_layer(inputs.view(inputs.shape[0], -1))
    loss = loss_fn(logits, targets)
    loss.backward()

    vec_grad = torch.flatten(fc_layer.weight.grad)
    precond_adagrad = torch.outer(vec_grad, vec_grad)

    # CPU computation works fine
    evals_adagrad, evecs_adagrad = torch.linalg.eigh(precond_adagrad.cpu())

    # But eigh computation on GPU fails
    evals_adagrad, evecs_adagrad = torch.linalg.eigh(precond_adagrad)

Error trace

Traceback (most recent call last):
  File "bug_report.py", line 47, in <module>
    evals_adagrad, evecs_adagrad = torch.linalg.eigh(precond_adagrad)
torch._C._LinAlgError: cusolver error: CUSOLVER_STATUS_EXECUTION_FAILED, when calling `cusolverDnXsyevd( handle, params, jobz, uplo, n, CUDA_R_32F, reinterpret_cast<void*>(A), lda, CUDA_R_32F, reinterpret_cast<void*>(W), CUDA_R_32F, reinterpret_cast<void*>(bufferOnDevice), workspaceInBytesOnDevice, reinterpret_cast<void*>(bufferOnHost), workspaceInBytesOnHost, info)`. This error may appear if the input matrix contains NaN.

cc @ezyang @gchanan @zou3519 @ptrblck @jianyuh @nikitaved @pearu @mruberry @walterddr @IvanYashchuk @xwang233 @lezcano @ngimel @hjmshi @mikerabbat @tsunghsienlee @dmudiger

Versions

PyTorch version: 2.0.0.dev20230201+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.25.2
Libc version: glibc-2.31

Python version: 3.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.0-124-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: Quadro GP100
GPU 1: Quadro GP100

Nvidia driver version: 470.141.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   46 bits physical, 48 bits virtual
CPU(s):                          80
On-line CPU(s) list:             0-79
Thread(s) per core:              2
Core(s) per socket:              20
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           79
Model name:                      Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz
Stepping:                        1
CPU MHz:                         2414.868
CPU max MHz:                     3600.0000
CPU min MHz:                     1200.0000
BogoMIPS:                        4399.78
Virtualization:                  VT-x
L1d cache:                       1.3 MiB
L1i cache:                       1.3 MiB
L2 cache:                        10 MiB
L3 cache:                        100 MiB
NUMA node0 CPU(s):               0-19,40-59
NUMA node1 CPU(s):               20-39,60-79
Vulnerability Itlb multihit:     KVM: Mitigation: Split huge pages
Vulnerability L1tf:              Mitigation; PTE Inversion; VMX conditional cache flushes, SMT vulnerable
Vulnerability Mds:               Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Meltdown:          Mitigation; PTI
Vulnerability Mmio stale data:   Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Mitigation; Clear CPU buffers; SMT vulnerable
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cdp_l3 invpcid_single pti intel_ppin ssbd ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm cqm rdt_a rdseed adx smap intel_pt xsaveopt cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts md_clear flush_l1d

Versions of relevant libraries:
[pip3] mypy==0.991
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.24.0
[pip3] pytorch-triton==2.0.0+0d7e753227
[pip3] torch==2.0.0.dev20230201+cu118
[pip3] torchvision==0.15.0.dev20230201+cu118
[conda] numpy                     1.24.0                   pypi_0    pypi
[conda] pytorch-triton            2.0.0+0d7e753227          pypi_0    pypi
[conda] torch                     2.0.0.dev20230201+cu118          pypi_0    pypi
[conda] torchvision               0.15.0.dev20230201+cu118          pypi_0    pypi

[hjmshi]

Just want to emphasize that this is a major issue for experimenting with full-matrix Adagrad and Shampoo. Any suggestions on how to proceed would be much appreciated!

cc: @shintaro-iwasaki @0x10cxR1 @kaushik88 @mhfb22 @dmudigere

[dmudigere]

also cc: @xwang233

[IvanYashchuk]

The matrix is all zeros and unfortunately cuSOLVER backend doesn't handle this edge case without an error. Could you try initializing vec_grad to non-zero values?

[gallego-posada]

If you are referring to precond_adagrad, this matrix is not all zero. It is the outer product of the gradient with itself, and the gradient is non-zero for at the initialization.

For a typical run I observe:

(Pdb) precond_adagrad.abs().max()
tensor(0.0697, device='cuda:0')
(Pdb) precond_adagrad.abs().min()
tensor(3.6154e-13, device='cuda:0')
(Pdb) precond_adagrad.abs().median()
tensor(0.0002, device='cuda:0')

[hjmshi]

Thanks for looking into it @IvanYashchuk! FYI, @mikerabbat has observed failure cases of torch.linalg.eigh on sparse diagonal matrices, with 6-7 non-zero entries on the diagonal. (This can be reproduced by running full-matrix Adagrad or Shampoo on a dense embedding table for sequence-to-sequence modeling, for example.) Should we consider incorporating checks on the matrix for these special cases, either in PyTorch or cuSOLVER? Thanks again.

[lezcano]

It'd be good to have the exact matrix for which this one fails. Could you guys provide that example?

My money is on the fact that the input matrix has repeated eigenvalues. Eigenvalue solvers struggle a bit with repeated eigenvalues...

[gallego-posada]

Thanks @lezcano -- I have saved the numpy array for to thevec_grad tensor here:
vec_grad.txt

You can restore the matrix precond_adagrad by:

import torch
import numpy as np

vec_grad = torch.tensor(np.loadtxt("vec_grad.out")).float()
precond_adagrad = torch.outer(vec_grad, vec_grad)

[lezcano]

Right. An outer product gives a matrix of rank one. As such, in this case, you have a matrix of shape 7840x7840 with 1 non-zero eigenvalue and 7839 eigenvalues equal to zero. Eigensolvers choke on these, as their performance often depends on the inverse of the eigengap of the matrix.

In other words, your solution to the problem is a bit numerically unstable, similar to when you try to solve a linear system with a matrix that's almost singular.

What I would suggest is that, if you know that your matrix has the rank-1 structure, or low-rank structure in general, you use other methods. For a rank-1 symmetric matrix of the form vv^T it's easy. Its non-zero eigenvalue is simply given by the square of the norm of v. For low-rank matrices, there are iterative methods in literature that are fast and stable if you have a bound on the rank of the matrix (which I believe you do in your problem). Alas, we don't have any of these methods in PyTorch.

[hjmshi]

@lezcano, thanks for your response. Just as additional context: torch.linalg.eigh does not fail on these matrices on CPU, this is specific to GPU/CUDA. We have observed this problem with CUDA 11.6, 11.7, and 11.8, but we do not observe these failures with CUDA 11.4. This also may be resolved in CUDA 12, but it would be great if we can have these problems ironed out on earlier versions since the current stable build of PyTorch suggests CUDA 11.6 and 11.7.

We will try to dig up some additional failure cases where the matrix is not just low-rank and post our findings as we also investigate further on our side. Thanks!

[lezcano]

Note that the difference of behaviour between CPU and CUDA and this exact point on matrices with repeated eigenvalues is documented here: https://pytorch.org/docs/master/notes/numerical_accuracy.html#extremal-values-in-linalg

[hjmshi]

Hi @lezcano, just want to clarify that on CPU, the algorithm at least succeeds (perhaps with numerical error), while on GPU, the algorithm entirely fails. We are willing to tolerate numerical error in the eigenvectors if the matrix is low-rank, as it will still produce an orthonormal basis for the zero eigenvalue subspace. I am referring here to the error bounds described here: https://netlib.org/lapack/lug/node90.html. Based on my understanding of your argument, even if we regularize by some constant here by adding epsilon * I, wouldn't the algorithm still fail? Would it even fail on a diagonal matrix c * I?

If you recommend using a different algorithm for the low rank case, which algorithms would you suggest? Thanks again.

[lezcano]

The error bounds in the link you provide say what I mentioned above. You see that $\theta$ is bounded by a expression that depends on the inverse of the eigengap so...
You'd have the same issue if you add eps * I as the eigenvalues of eps * I + vv^T are n-1 equal to eps and one equal to eps + \norm{v}^2.

Theoretically, these algorithms may even fail for diagonal matrices, but in this case, since the matrix is exactly low-rank on floating point precision, the algorithm often succeeds. Note that vv^T is just 1-rank in floating point precision, but when you perform operations on it you may get some very small numbers that which cause these algorithms to fail.

As to which algorithm to use... well, that very much depends on your exact problem. I'd recommend you discuss it with your local linear algebra expert. For large and very sparse matrices, you have the classic SVD approximation described in https://arxiv.org/abs/0909.4061, that I believe would work well for your case, as you can give a bound on the rank of your matrix. You even have a post from meta describing what's what, and there is a fast implementation in cusolver of it. Again, different algorithms may be more amenable to different problems.

[anana10c]

Thanks @lezcano for the quick response on #105359! I understand that failure is expected on nearly singular matrices, but the matrix becoming unrecoverable afterwards may be an issue for our use case. Is there anything that can be done besides implementing checks on our side to convert to float64 before calling eigh, or preconditioning the matrix as you suggested?

Additionally, I was able to reproduce the error yesterday on CUDA 12.1 as well. It seems to only occur on matrices larger than 512 x 512 - is the issue likely to be in the syevd solver?

[lezcano]

The issue is indeed in the cusolver side, so there's very little we can do on our end really. I don't think that there's any other solution for this ATM.

[hjmshi]

cc: @dmudigere @xwang233

[xwang233]

Thanks for the reminder. We'll take a look

[gottbrath]

What is it that is left in an unrecoverable state? The tensor? The GPU?

[anana10c]

It seems that nothing on the GPU can be accessed anymore without a CUDA illegal memory access error.

import torch

factor_matrix = torch.load("rank7_idx0.1.3.0_iter100_factor.pt")
factor_matrix = factor_matrix.to("cuda:0")

tmp = torch.arange(10)
tmp = tmp.to("cuda:0")

torch.linalg.eigh(factor_matrix)  # fails

print(tmp)  # CUDA illegal memory access

It turns out that I may have been incorrect about the error occurring in CUDA 12.1 - after changing some environment variables, it seems to no longer fail? I'll investigate this some more and keep everyone updated.

[anana10c]

Seems like I was running on PyTorch nightly (2.1.0.dev20230717+cu121) by accident. Here's what I seem to have observed so far (please let me know if you can reproduce this!):

• PyTorch from source (2.1.0a0+git6855053) + CUDA 11.8 fails (causes unrecoverable state)
• PyTorch nightly (2.1.0.dev20230717+cu121) + CUDA 12.1 fails (causes unrecoverable state)
• PyTorch from source (2.1.0a0+git6855053) + CUDA 12.1 succeeds (no cusolver error)

If another matrix were to fail in the last setup, I'm unsure whether it would cause the unrecoverable state or not.

[mrogowski]

I believe this issue has been fixed starting from cuSOLVER 11.4.5, which shipped with CUDA 12.1 Update 1.

Can you please confirm if this matches your observations?

[anana10c]

As before, while everything seems fine with PyTorch built from source + CUDA 12.1, if I use PyTorch nightly + CUDA 12.1 instead the error still occurs:

Python 3.9.17 (main, Jul  5 2023, 20:41:20) 
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.__version__
'2.1.0.dev20230721+cu121'
>>> torch.version.cuda
'12.1'
>>> factor_matrix = torch.load("rank7_idx0.1.3.0_iter100_factor.pt")
>>> factor_matrix = factor_matrix.to("cuda:0")
>>> torch.linalg.eigh(factor_matrix)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: cusolver error: CUSOLVER_STATUS_EXECUTION_FAILED, when calling `cusolverDnXsyevd( handle, params, jobz, uplo, n, CUDA_R_32F, reinterpret_cast<void*>(A), lda, CUDA_R_32F, reinterpret_cast<void*>(W), CUDA_R_32F, reinterpret_cast<void*>(bufferOnDevice), workspaceInBytesOnDevice, reinterpret_cast<void*>(bufferOnHost), workspaceInBytesOnHost, info)`
>>> factor_matrix
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/fsx/users/annacai/conda/envs/eigh2/lib/python3.9/site-packages/torch/_tensor.py", line 430, in __repr__
    return torch._tensor_str._str(self, tensor_contents=tensor_contents)
  File "/fsx/users/annacai/conda/envs/eigh2/lib/python3.9/site-packages/torch/_tensor_str.py", line 674, in _str
    return _str_intern(self, tensor_contents=tensor_contents)
  File "/fsx/users/annacai/conda/envs/eigh2/lib/python3.9/site-packages/torch/_tensor_str.py", line 605, in _str_intern
    tensor_str = _tensor_str(self, indent)
  File "/fsx/users/annacai/conda/envs/eigh2/lib/python3.9/site-packages/torch/_tensor_str.py", line 357, in _tensor_str
    formatter = _Formatter(get_summarized_data(self) if summarize else self)
  File "/fsx/users/annacai/conda/envs/eigh2/lib/python3.9/site-packages/torch/_tensor_str.py", line 393, in get_summarized_data
    return torch.stack([get_summarized_data(x) for x in (start + end)])
  File "/fsx/users/annacai/conda/envs/eigh2/lib/python3.9/site-packages/torch/_tensor_str.py", line 393, in <listcomp>
    return torch.stack([get_summarized_data(x) for x in (start + end)])
  File "/fsx/users/annacai/conda/envs/eigh2/lib/python3.9/site-packages/torch/_tensor_str.py", line 383, in get_summarized_data
    return torch.cat(
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

[anana10c]

Is there a way for me to check the cusolver version? My CUDA runtime version seems to be 12.1.105.

[mrogowski]

You can check the version stated in the filename of the library, for example: /usr/local/cuda/lib64/libcusolver.so.11.4.5.107

[anana10c]

I see, libcusolver.so.11.4.5.107 seems to exist in /usr/local/cuda-12.1/lib64 for me. Both versions of PyTorch that I've tried should be built on the same CUDA version (in /usr/local/cuda-12.1).

[mrogowski]

I was able to reproduce the issue using the script from the first post, but that error was definitely resolved in CUDA 12.1 Update 1. Can you share your file, rank7_idx0.1.3.0_iter100_factor.pt, so that I can try to reproduce your issue?

Also, please note that cuSOLVER had this issue in CUDA 12.1, and it was fixed in the update. You can check the version with nvcc -V (you will see 12.1.66 for CUDA 12.1 and 12.1.105 for CUDA 12.1 Update 1).

[anana10c]

The file rank7_idx0.1.3.0_iter100_factor.pt can be found here: #105359

nvcc -V returns:

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Mon_Apr__3_17:16:06_PDT_2023
Cuda compilation tools, release 12.1, V12.1.105
Build cuda_12.1.r12.1/compiler.32688072_0

[mrogowski]

Regarding the CUDA version discussion, this is fixed in 12.1 Update 1.

[malfet]

@lezcano I'm totally fine with numerical errors, but imo any types of crashes/unrecoverable errors on a valid input should be considered as high-pri (as it breaks the pipeline). Also, based on @anana10c comment, added module: regression to the list, as the same used to work with CUDA-11.4, but is broken with CUDA-11.8 and above.

Regarding: "there is not much else we can do but wait", I disagree with you: we can stick to CUDA-11.4, we can use magma and so on.

[soulitzer]

From discussion in triage review:

• we should add a test to prevent regressions
• properly document support wrt different CUDA versions
• possibly add support using MAGMA

[xwang233]

We'll add a test case for those inputs.

For "possibly add support using MAGMA", users may refer to this function to prefer MAGMA as the linear algebra backend library. This is a global runtime switch. https://pytorch.org/docs/stable/backends.html#torch.backends.cuda.preferred_linalg_library

torch.backends.cuda.preferred_linalg_library('magma')

There might be performance hits when switching from default (cusolver) to MAGMA in eigh and other eigh-based functions.

[ptrblck]

Reproduced the issue using 2.0.1+cu117 as well as a recent nightly with 11.8: torch-2.1.0.dev20230823+cu118. The latest 12.1Update1 nightly works fine and does not reproduce the issue anymore: torch-2.1.0.dev20230823+cu121.