Error with jnp.linalg.det method on cuda device

[unik-w]

jnp.linalg.det(array) method raises error when run on a cuda device. When run using a cpu instead the determinant is calculated without any error.

The following code describes the reproduction of the issue:

import os

os.environ['JAX_PLATFORMS'] = 'cuda'
import jax
print(jax.devices())

import jax.numpy as jnp
from jax import random



def det():
    key = random.PRNGKey(0) 
    array = random.normal(key, shape=(10, 10))  
    return jnp.linalg.det(array)

det()

Error:  Cell In[4], <a href='vscode-notebook-cell:?execution_count=4&line=16'>line 16</a>
     <a href='vscode-notebook-cell:?execution_count=4&line=14'>14</a> key = random.PRNGKey(0) 
     <a href='vscode-notebook-cell:?execution_count=4&line=15'>15</a> array = random.normal(key, shape=(10, 10))  
---> <a href='vscode-notebook-cell:?execution_count=4&line=16'>16</a> return jnp.linalg.det(array)

    [... skipping hidden 33 frame]

File [~/MasterArbeit/tc-nqs/.venvjax/lib/python3.11/site-packages/jaxlib/gpu_solver.py:102](https://vscode-remote+jupyter2-002ephysik-002euni-002dmuenchen-002ede.vscode-resource.vscode-cdn.net/home/u/Unik.Wadhwani/MasterArbeit/tc-nqs/notes/~/MasterArbeit/tc-nqs/.venvjax/lib/python3.11/site-packages/jaxlib/gpu_solver.py:102), in _getrf_hlo(platform, gpu_blas, gpu_solver, dtype, a)
    <a href='~/MasterArbeit/tc-nqs/.venvjax/lib/python3.11/site-packages/jaxlib/gpu_solver.py:100'>100</a>   kernel = f"{platform}blas_getrf_batched"
    <a href='~/MasterArbeit/tc-nqs/.venvjax/lib/python3.11/site-packages/jaxlib/gpu_solver.py:101'>101</a> else:
--> <a href='~/MasterArbeit/tc-nqs/.venvjax/lib/python3.11/site-packages/jaxlib/gpu_solver.py:102'>102</a>   lwork, opaque = gpu_solver.build_getrf_descriptor(
    <a href='~/MasterArbeit/tc-nqs/.venvjax/lib/python3.11/site-packages/jaxlib/gpu_solver.py:103'>103</a>       np.dtype(dtype), batch, m, n)
    <a href='~/MasterArbeit/tc-nqs/.venvjax/lib/python3.11/site-packages/jaxlib/gpu_solver.py:104'>104</a>   workspace = ir.RankedTensorType.get([lwork], a_type.element_type)
    <a href='~/MasterArbeit/tc-nqs/.venvjax/lib/python3.11/site-packages/jaxlib/gpu_solver.py:105'>105</a>   kernel = f"{platform}solver_getrf"

RuntimeError: jaxlib/gpu/solver_kernels.cc:45: operation gpusolverDnCreate(&handle) failed: cuSolver internal error

Jax and jaxlib version: 0.4.28
cuda driver 12.4

[dfm]

This looks like an issue with your JAX installation. Can you try installing the most recent version:

pip install -U "jax[cuda12]"

[hawkinsp]

My other guess would be you're low on GPU memory. Can you also include the output of nvidia-smi?

[unik-w]

Thanks for the reply
I updates the jax and jaxlib version to 0.4.31 using pip install -U "jax[cuda12]". that doesnt help with the problem.
All other jax.numpy methods(and more jax methods such as vmap/jit/grad) run normally on cuda device (other than jax.numpy.linalg.)

The above code for reproduction is just a small example from a larger cod based on jax methods.

[unik-w]

@hawkinsp Fri Aug 16 14:43:41 2024
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 555.42.06 Driver Version: 555.42.06 CUDA Version: 12.5 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 Quadro P2000 On | 00000000:61:00.0 Off | N/A |
| 49% 37C P8 4W / 75W | 3851MiB / 5120MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
+-----------------------------------------------------------------------------------------+

[hawkinsp]

Try lowering the JAX GPU memory fraction to something like 0.6 (https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html)

[unik-w]

@hawkinsp @dfm . Thanks for your answers.
However, The solutions provided do not work.
For context: I have a flax neural network with a step involving taking a determinant of a large matrix(using jnp.linalg.det or jnp.linalg.slogdet) , The rest of the network trains on GPU if i remove the determinant part or replace it with some other functions. (tested with Nvidia based GPUs Quadro P2000, A40,A100 gpus)
temperory solution: To use "jax.jit(jnp.linalg.det,device=jax.devices('cpu')[0])(x)" to explicitly calculate the determinant on CPU. and rest of the process on GPU.
But this is not very efficient solution for me.