torch.linalg.lstsq gives wrong result on GPU

[fuzzyswan]

🐛 Describe the bug

torch.linalg.lstsq on cpu and gpu gives different results, and the gpu result is different from numpy.

import torch
import numpy as np

a = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
b = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
c = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]).cuda()
d = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]).cuda()
e = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]).detach().numpy()
f = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]).detach().numpy()

x = torch.linalg.lstsq(a, b)[0]
y = torch.linalg.lstsq(c, d)[0]
z = np.linalg.lstsq(e, f)[0]

print(x)
print(y)
print(z)

Results:

tensor([[ 0.8333,  0.3333, -0.1667],
        [ 0.3333,  0.3333,  0.3333],
        [-0.1667,  0.3333,  0.8333]])
tensor([[ 0.5814,  0.7631, -0.0000],
        [ 0.8372, -0.5261,  0.0000],
        [-0.4186,  0.7631,  1.0000]], device='cuda:0')
[[ 0.8333333   0.33333334 -0.16666667]
 [ 0.33333334  0.33333334  0.33333334]
 [-0.16666667  0.33333334  0.8333333 ]]

Versions

Collecting environment information...
PyTorch version: 1.12.1+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.22.6
Libc version: glibc-2.26

Python version: 3.7.15 (default, Oct 12 2022, 19:14:55) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.10.133+-x86_64-with-Ubuntu-18.04-bionic
Is CUDA available: True
CUDA runtime version: 11.2.152
CUDA_MODULE_LOADING set to:
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 460.32.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.1.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.1.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.6
[pip3] torch==1.12.1+cu113
[pip3] torchaudio==0.12.1+cu113
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.13.1
[pip3] torchvision==0.13.1+cu113
[conda] Could not collect

cc @ngimel @jianyuh @nikitaved @pearu @mruberry @walterddr @IvanYashchuk @xwang233 @lezcano

[ezyang]

your test matrix is not well conditioned

>>> import torch
>>> a = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
>>> torch.linalg.svd(a)
torch.return_types.linalg_svd(
U=tensor([[-0.2148,  0.8872,  0.4082],
        [-0.5206,  0.2496, -0.8165],
        [-0.8263, -0.3879,  0.4082]]),
S=tensor([1.6848e+01, 1.0684e+00, 1.0313e-07]),
Vh=tensor([[-0.4797, -0.5724, -0.6651],
        [-0.7767, -0.0757,  0.6253],
        [-0.4082,  0.8165, -0.4082]]))

[IvanYashchuk]

torch.linalg.lstsq has different default algorithms on CPU and CUDA. CPU behavior matches NumPy because the same algorithm is used.

[lezcano]

To elaborate on @ezyang's point, as it says in the docs, the CUDA implementation just supports matrices that are full-rank.

[fuzzyswan]

Thank you! I have just re-checked the docs and it is clear on this.