'false INTERNAL ASSERT FAILED' when calling torch.pinverse and using torch.float32 on Apple Silicon

[PlayerCham]

🐛 Describe the bug

When running the torch.pinverse function, the ‘false INTERNAL ASSERT FAILED’ error only occurs when running on an Apple Silicon (mine is M3 Pro) device using torch.float32, and not when using torch.float64. When running on an X86_64 (mine is AMD Ryzen 9 5950X) device, the error does not occur when using torch.float32 or 64.

The error will occur very quickly, otherwise it will run normally for a long time until the calculation is completed.

import torch

def reproduce_bug():
    m, n, k = 50000, 8193, 10
    H_aug = torch.randn(n, m, dtype=torch.float32)  #Change to torch.float64 here
    T = torch.randn(n, k, dtype=torch.float32)  #Change to torch.float64 here
    W_aug = torch.pinverse(H_aug) @ T
    print("W_aug shape:", W_aug.shape)

if __name__ == "__main__":
    reproduce_bug()

(test) yuyao@cyys-MacBook-Pro-14-M3-Pro Downloads % python float32.py 
** On entry to SGESDD, parameter number 12 had an illegal value
Traceback (most recent call last):
  File "/Users/yuyao/Downloads/float32.py", line 11, in <module>
    reproduce_bug()
  File "/Users/yuyao/Downloads/float32.py", line 7, in reproduce_bug
    W_aug = torch.pinverse(H_aug) @ T
            ^^^^^^^^^^^^^^^^^^^^^
RuntimeError: false INTERNAL ASSERT FAILED at "/Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/BatchLinearAlgebra.cpp":1602, please report a bug to PyTorch. linalg.svd: Argument 12 has illegal value. Most certainly there is a bug in the implementation calling the backend library.

Versions

(test) yuyao@cyys-MacBook-Pro-14-M3-Pro Downloads % python collect_env.py 
Collecting environment information...
PyTorch version: 2.6.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 15.3.2 (arm64)
GCC version: Could not collect
Clang version: 16.0.0 (clang-1600.0.26.6)
CMake version: Could not collect
Libc version: N/A

Python version: 3.12.9 | packaged by Anaconda, Inc. | (main, Feb  6 2025, 12:55:12) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-15.3.2-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M3 Pro

Versions of relevant libraries:
[pip3] numpy==2.2.4
[pip3] torch==2.6.0
[pip3] torchvision==0.21.0
[conda] numpy                     2.2.4                    pypi_0    pypi
[conda] torch                     2.6.0                    pypi_0    pypi
[conda] torchvision               0.21.0                   pypi_0    pypi

(test) C:\Users\yuyao\Downloads>python collect_env.py
Collecting environment information...
PyTorch version: 2.6.0+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 专业版 (10.0.26100 64 位)
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.12.9 | packaged by Anaconda, Inc. | (main, Feb  6 2025, 18:49:16) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-11-10.0.26100-SP0
Is CUDA available: False
CUDA runtime version: 12.1.66
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3080 Ti
Nvidia driver version: 572.70
cuDNN version: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1\bin\cudnn_ops_train64_8.dll
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Name: AMD Ryzen 9 5950X 16-Core Processor
Manufacturer: AuthenticAMD
Family: 107
Architecture: 9
ProcessorType: 3
DeviceID: CPU0
CurrentClockSpeed: 3401
MaxClockSpeed: 3401
L2CacheSize: 8192
L2CacheSpeed: None
Revision: 8448

Versions of relevant libraries:
[pip3] numpy==2.2.4
[pip3] torch==2.6.0
[pip3] torchvision==0.21.0
[conda] numpy                     2.2.4                    pypi_0    pypi
[conda] torch                     2.6.0                    pypi_0    pypi
[conda] torchvision               0.21.0                   pypi_0    pypi

cc @malfet @albanD @jianyuh @nikitaved @pearu @mruberry @walterddr @xwang233 @lezcano @snadampal @milpuz01

[malfet]

Looks eerily reminiscent of #145801 (cc: @wdvr )
Indeed it is, grabbing for myself, just for fun, following patch fixes it:

diff --git a/aten/src/ATen/native/BatchLinearAlgebraKernel.cpp b/aten/src/ATen/native/BatchLinearAlgebraKernel.cpp
index 8dce552b0e1..e2e8a9e8918 100644
--- a/aten/src/ATen/native/BatchLinearAlgebraKernel.cpp
+++ b/aten/src/ATen/native/BatchLinearAlgebraKernel.cpp
@@ -1066,7 +1066,8 @@ static void apply_svd(const Tensor& A,
   {
     scalar_t wkopt;
     lapackSvd<scalar_t, value_t>(jobz, m, n, A_data, lda, S_data, U_data, ldu, Vh_data, ldvh, &wkopt, lwork, rwork_data, iwork_data, info_data);
-    lwork = std::max<int>(1, real_impl<scalar_t, value_t>(wkopt));
+    value_t next_after_lw = std::nextafter(real_impl<scalar_t, value_t>(wkopt), std::numeric_limits<value_t>::infinity());
+    lwork = std::max<int>(1, std::ceil(next_after_lw));
   }