cuSolver internal error when running simulations

[Zolkin1]

Occasionally I get this error:

Traceback (most recent call last):
  File "/hydrax/examples/pusht.py", line 34, in <module>
    run_interactive(
  File "/hydrax/hydrax/simulation/deterministic.py", line 177, in run_interactive
    policy_params, rollouts = jit_optimize(mjx_data, policy_params)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/pjit.py", line 339, in cache_miss
    pgle_profiler) = _python_pjit_helper(fun, jit_info, *args, **kwargs)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/pjit.py", line 194, in _python_pjit_helper
    out_flat, compiled, profiler = _pjit_call_impl_python(*args_flat, **p.params)
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/pjit.py", line 1659, in _pjit_call_impl_python
    ).compile()
      ^^^^^^^^^
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 2448, in compile
    executable = UnloadedMeshExecutable.from_hlo(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 2967, in from_hlo
    xla_executable = _cached_compilation(
                     ^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 2758, in _cached_compilation
    xla_executable = compiler.compile_or_get_cached(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/compiler.py", line 470, in compile_or_get_cached
    return _compile_and_write_cache(
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/compiler.py", line 687, in _compile_and_write_cache
    executable = backend_compile(
                 ^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/profiler.py", line 334, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/compiler.py", line 327, in backend_compile
    raise e
  File "/miniconda3/envs/hydrax/lib/python3.12/site-packages/jax/_src/compiler.py", line 321, in backend_compile
    return backend.compile(built_c, compile_options=options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: cuSolver internal error
Xlib:  extension "NV-GLX" missing on display ":1".

(this is run with JAX_TRACEBACK_FILTERING=off)

I believe others have got this error too. It seems to be very hit or miss - sometimes it happens a few times in a row and sometimes not at all.

Not sure if anyone has any insight here, but it can be annoying. I might be able to put a few cycles into chasing it down, but not sure.

[Tom0Brien]

I suspect it's something to do with trying to allocate to much gpu memory jax-ml/jax#8916.

For example, you can reproduce this problem by trying to spin up multiple simulations at once.

Setting the following looks like a potential solution to help:
export XLA_PYTHON_CLIENT_PREALLOCATE=false

[vincekurtz]

Maybe XLA_PYTHON_CLIENT_MEM_FRACTION could help as well? But it is strange to me that this would throw a cuSolver error rather than an OOM error

[itaouil]

Hi,

I recently updated my nvidia driver on my laptop (RTX 4060), and since then I have not been able to launch the script. In fact, every single time I launch it I receive the following error:

raceback (most recent call last): File "/home/ilyass/workspace/hydrax/examples/humanoid_mocap.py", line 72, in <module> run_interactive( File "/home/ilyass/workspace/hydrax/hydrax/simulation/deterministic.py", line 96, in run_interactive policy_params, rollouts = jit_optimize(mjx_data, policy_params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 182, in reraise_with_filtered_traceback return fun(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^ File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/pjit.py", line 334, in cache_miss executable, pgle_profiler) = _python_pjit_helper(fun, jit_info, *args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/pjit.py", line 195, in _python_pjit_helper out_flat, compiled, profiler = _pjit_call_impl_python(*args_flat, **p.params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/pjit.py", line 1862, in _pjit_call_impl_python ).compile() ^^^^^^^^^ File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 2467, in compile executable = UnloadedMeshExecutable.from_hlo( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 3009, in from_hlo xla_executable = _cached_compilation( ^^^^^^^^^^^^^^^^^^^^ File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 2800, in _cached_compilation xla_executable = compiler.compile_or_get_cached( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/compiler.py", line 447, in compile_or_get_cached return _compile_and_write_cache( ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/compiler.py", line 719, in _compile_and_write_cache executable = backend_compile( ^^^^^^^^^^^^^^^^ File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/profiler.py", line 354, in wrapper return func(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/compiler.py", line 342, in backend_compile raise e File "/home/ilyass/mambaforge/envs/hydrax/lib/python3.12/site-packages/jax/_src/compiler.py", line 335, in backend_compile return backend.compile( ^^^^^^^^^^^^^^^^ jaxlib._jax.XlaRuntimeError: INTERNAL: cuSolver internal error

This is my JAX environment info:

`jax: 0.6.1
jaxlib: 0.6.1
numpy: 2.2.6
python: 3.12.10 | packaged by conda-forge | (main, Apr 10 2025, 22:21:13) [GCC 13.3.0]
device info: NVIDIA GeForce RTX 4060 Laptop GPU-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='atari', release='6.8.0-59-generic', version='#61~22.04.1-Ubuntu SMP PREEMPT_DYNAMIC Tue Apr 15 17:03:15 UTC 2', machine='x86_64')
JAX_TRACEBACK_FILTERING=off

$ nvidia-smi
Fri May 23 22:17:29 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.144.03 Driver Version: 550.144.03 CUDA Version: 12.4 |
|-----------------------------------------+------------------------+----------------------+
| 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 NVIDIA GeForce RTX 4060 ... Off | 00000000:01:00.0 Off | N/A |
| N/A 41C P3 2W / 50W | 112MiB / 8188MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| 0 N/A N/A 3153 G /usr/lib/xorg/Xorg 4MiB |
| 0 N/A N/A 20374 C python 98MiB |
+-----------------------------------------------------------------------------------------+
`

and this is my nvidia driver:

nvcc: NVIDIA (R) Cuda compiler driver Copyright (c) 2005-2023 NVIDIA Corporation Built on Fri_Sep__8_19:17:24_PDT_2023 Cuda compilation tools, release 12.3, V12.3.52 Build cuda_12.3.r12.3/compiler.33281558_0

have you ever experienced something like this?
Could you share your conda/nvidia setup (driver, cuda, cudnn) to try and see if the error is maybe due to some library mismatch?

For completness I also attach my conda environmet as a txt file so you can see the version of my libraries:

environment_ilyass.txt

Thanks :)

[vincekurtz]

@itaouil My env has the following:

$ pip list | grep nvidia
nvidia-cublas-cu12       12.6.3.3
nvidia-cuda-cupti-cu12   12.6.80
nvidia-cuda-nvcc-cu12    12.6.77
nvidia-cuda-runtime-cu12 12.6.77
nvidia-cudnn-cu12        9.5.1.17
nvidia-cufft-cu12        11.3.0.4
nvidia-cusolver-cu12     11.7.1.2
nvidia-cusparse-cu12     12.5.4.2
nvidia-nccl-cu12         2.23.4
nvidia-nvjitlink-cu12    12.6.77

$ pip list | grep jax
jax                      0.4.35
jax-cuda12-pjrt          0.4.35
jax-cuda12-plugin        0.4.35
jaxlib                   0.4.34
jaxtyping                0.3.1

And from nvidia-smi: Driver Version: 550.144.03 CUDA Version: 12.4

[itaouil]

Thanks.

I managed to solve the issue. The way I did it is to install the cuda driver, cudnn, and cuda toolkit myself to match the jax version that I wanted, and then only installing the "local" version of jax using pip install jax["cuda12-local"] otherwise jax installs nvidia libraries that do not match in my case.

[giorgionicola]

Hi, I have experienced a similar error but somehow weird. First of all at the beginning was happening with al examples and everything and I solvde it bz updating nvidia-cublas-cu12==12.9.0.13 based on https://github.com/jax-ml/jax/issues/29042.

That fixed most of my problems. However when i try to run the example in asynchronous mode it gives the same error.

Running asynchronous simulation
Planning with 30 steps over a 0.6 second horizon with 4 knots.
Jitting controller...
Process SpawnProcess-2:
Traceback (most recent call last):
  File "/home/giorgionicola/anaconda3/envs/hydrax/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/home/giorgionicola/anaconda3/envs/hydrax/lib/python3.12/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/giorgionicola/lib/hydrax/hydrax/simulation/asynchronous.py", line 134, in run_controller
    policy_params = jit_optimize(mjx_data, policy_params)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib._jax.XlaRuntimeError: INTERNAL: cuSolver internal error
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

my pip list | grep nvidia

nvidia-cublas-cu12       12.9.0.13
nvidia-cuda-cupti-cu12   12.9.79
nvidia-cuda-nvcc-cu12    12.9.86
nvidia-cuda-runtime-cu12 12.9.79
nvidia-cudnn-cu12        9.10.2.21
nvidia-cufft-cu12        11.4.1.4
nvidia-cusolver-cu12     11.7.5.82
nvidia-cusparse-cu12     12.5.10.65
nvidia-nccl-cu12         2.27.3
nvidia-nvjitlink-cu12    12.9.86
nvidia-nvshmem-cu12      3.2.5