DESC installation on Perlmutter doesn't work

[rahulgaur104]

I followed the installation instruction verbatim and I get this on Perlmutter.

module load cudatoolkit/12.4
module load cudnn/9.5.0
module load conda

conda create -n desc-env python=3.12
conda activate desc-env
pip install --upgrade "jax[cuda12]"

git clone https://github.com/PlasmaControl/DESC.git
cd DESC
pip install --editable .

To test my installation, I import Equilibrium from my login node

(desc-env) rgaur@login09:/pscratch/sd/r/rgaur/DESC> python3
Python 3.12.11 | packaged by Anaconda, Inc. | (main, Jun  5 2025, 13:09:17) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
Ctrl click to launch VS Code Native REPL
>>> import desc
>>> from desc.equilibrium import Equilibrium
ERROR:2025-09-19 06:12:41,224:jax._src.xla_bridge:444: Jax plugin configuration error: Exception when calling jax_plugins.xla_cuda12.initialize()
Traceback (most recent call last):
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 201, in _version_check
    version = get_version()
              ^^^^^^^^^^^^^
RuntimeError: jaxlib/cuda/versions_helpers.cc:81: operation cusparseGetProperty(MAJOR_VERSION, &major) failed: The cuSPARSE library was not found.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax/_src/xla_bridge.py", line 442, in discover_pjrt_plugins
    plugin_module.initialize()
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 328, in initialize
    _check_cuda_versions(raise_on_first_error=True)
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 266, in _check_cuda_versions
    _version_check("cuSPARSE", cuda_versions.cusparse_get_version,
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 205, in _version_check
    raise RuntimeError(err_msg) from e
RuntimeError: Unable to load cuSPARSE. Is it installed?

Adding set_device("gpu") before importing Equilibrium doesn't help either.

Also, by default, this process installated jax = jaxlib = 0.6.2 while jax-cuda12-pjrt = jax-cuda12-plugin = 0.7.2.
Regardless of the version, DESC installation doesn't work on Perlmutter.

[f0uriest]

I tested with jax 0.6.2, 0.7.0 and 0.7.1 over last weekend and seemed fine. It's possible 0.7.2 (released this week) broke something or needs a different cuda version? But I don't see anything in the changelog. And you're saying 0.6.2 also doesn't work?

[f0uriest]

OK so I think there's several things going on here

1. Our requirements.txt pins jax <= 0.6.2 but pip install --upgrade "jax[cuda12]" installs a newer version, which then gets partially overwritten by pip install -r requirements.txt. We should fix this.
2. If I switch to doing pip install --upgrade "jax[cuda12]==0.6.2" I still run into similar issues, even though jax is now fully self consistent. The fix seems to be to NOT do module load cudatoolkit cudnn etc. I think this is no longer necessary and is causing issues with the version of cuda that jax pulls.
3. If I module load cuda... and then try to use an environment I already have installed with jax, that worked a few days ago, I get the same error as above.
4. If I use a jupyter notebook with the old environment, it seems to work fine - note that the jupyter server doesn't do any module load cuda ...
5. If I do

module purge
module load conda
conda create -n test-env python=3.12
conda activate test-env
pip install --upgrade "jax[cuda12]==0.6.2"
cd DESC
pip install --editable .

things seem to work fine.

so TLDR:

• Remove the module load cuda/cudnn etc from the perlmutter instructions
• make sure that the jax[cuda] version we install is compatible with the versions in requirements
• We should probably also bump the max version in requirements, not sure why dependabot has been missing this.

[rahulgaur104]

I was literally writing this response and I don't want to delete it. Some of this points have already been covered by @f0uriest above

I tested with jax 0.6.2, 0.7.0 and 0.7.1 over last weekend and seemed fine. It's possible 0.7.2 (released this week) broke something or needs a different cuda version? But I don't see anything in the changelog. And you're saying 0.6.2 also doesn't work?

I don't understand how what you say is possible. Perhaps, I should follow these instructions in front of you to show the problem.

Moreover, I realized that these installation instruction only use CPUs. With set_device("gpu"), I get

"RuntimeError: Unknown backend: 'gpu' requested, but no platforms that are instances of gpu are present. Platforms present are: cpu"

which means that it never installed gpu compatible versions of jax or cuda. I have checked my module list and it does contain the right version of cudatoolkit and cudnn.

I tested with jax 0.6.2, 0.7.0 and 0.7.1 over last weekend and seemed fine. It's possible 0.7.2 (released this week) broke something or needs a different cuda version? But I don't see anything in the changelog. And you're saying 0.6.2 also doesn't work?

@f0uriest, did you actually run an optimization problem on a GPU and make sure that the GPU is being used? What problem did you run?

As far as different versions are concerned, I don't understand how you were able to maintain

jax-cuda12-pjrt           0.6.2                    pypi_0    pypi
jax-cuda12-plugin         0.6.2                    pypi_0    pypi

because if I do everything from scratch, by default it should switch to the latest version which has been > 0.6.2 for a while (months) now but jax and jaxlib cannot go beyond 0.6.2 because the requirements.txt clips it to 0.6.2. I had to explicity install jax-cuda12-pjrt plugin == 0.6.2 to maintain the same version.

If someone were to follow the instructions exactly they would have obtained different version of jax/jaxlib and jax-cuda12-plugin/jax-cuda12-pjrt. However, I don't think that's the reason for the error.

These instructions never worked for me.
The problem arises from the fact that developers are not really using the code and users don't always report the bugs. We don't really work on how robust the code is. We should atleast make sure that it runs on the machined for which we provide installation instructions.
Perhaps, we should take some time to work on the robustness before any any new features or over-optimizing the code.

Here's what has been working for me on Perlmutter for years now

CONDA_OVERRIDE_CUDA="12.4" conda create --name desc-env "jax==0.6.0" "jaxlib==0.6.0=cuda12*" -c conda-forge
cd DESC
pip install --editable .
change the available jax version to the appropriate version. This can be obtained from conda-forge website.

[f0uriest]

I tested with jax 0.6.2, 0.7.0 and 0.7.1

By this I mean I installed different versions of jax[cuda12] in an existing environment and verified they work, but not creating a new environment or reinstalling DESC each time, so that may be the difference. The environment itself was made several months ago, presumably these instructions worked then? But I don't remember exactly what steps I followed (though I'm pretty sure I used whatever was on the docs at the time)

[YigitElma]

I think this line

Note that you may also need to execute unset LD_LIBRARY_PATH before starting a python process

was to prevent the cuSparse error. Did you try that?

I don't have access to Perlmutter, but other installation instructions call pip install "jax[cuda12]" later without --upgrade, that should resolve installing most recent/last supported jax version issue.

[f0uriest]

I have checked my module list and it does contain the right version of cudatoolkit and cudnn.

As of today I'm finding that loading any version of cudatoolkit (12.4 or 12.9) with any version of jax (0.6.2-0.7.2) causes errors. I suspect this is because jax includes some of its own cuda stuff, and if the module is loaded it grabs the wrong cuda (even though it appears to be the right version?). Which might also explain @YigitElma point about unset LD_LIBRARY_PATH

[rahulgaur104]

Note that you may also need to execute unset LD_LIBRARY_PATH before starting a python process

Ok, I have already spent ~ 2.5 hours installing, cleaning, setting-up and repeating this process on Perlmutter so I'll try it later but I don't think your suggestion about LD_LIBRARY_PATH will fix the problem. My guess is that this installation was never tested on a GPU so it was always failing.

Also, if it is necessary, why was it removed from the docs?

I don't have access to Perlmutter, but other installation instructions call pip install "jax[cuda12]" later without --upgrade, that should resolve installing most recent/last supported jax version issue.

Talk to Michael Churchill. Takes < 1 hour to get access. Please give me exact instructions if you want me to try them.

[YigitElma]

Also, if it is necessary, why was it removed from the docs?

It is still in the docs, not written as an explicit step because I think this was case-dependent (or JAX developers said it will be fixed, I don't remember exactly, this was the original post I found the fix at the time jax-ml/jax#28929). Also, this is not during installation but before running a JAX process, so you can technically use it with an existing environment.

(Again, I don't have access, so these are not tested) I would recommend using,

module load conda

conda create -n desc-env python=3.12
conda activate desc-env

cd DESC
pip install -r devtools/dev-requirements.txt   # or just requirements.txt
pip install "jax[cuda12]" # this already install all CUDA wheels

I personally omit using editable installs (I had a couple of problems before, but probably due to a stupid mistake of mine), so either use an editable install or this one. If you use this one, add

import sys
import os

sys.path.insert(0, os.path.abspath("."))
sys.path.append(os.path.abspath("../../../")) # where is your DESC/desc directory

to be able to add desc to your path. For unsetting the LD thing

os.environ["LD_LIBRARY_PATH "] = ""  # before any import from desc

or on command line (in slurm)

unset LD_LIBRARY_PATH 

before running a Python process.

[YigitElma]

Talk to Michael Churchill. Takes < 1 hour to get access.

Thanks, I will.

[rahulgaur104]

Here' what I did using the instructions provided by @YigitElma

module load conda

conda create -n desc-env python=3.12
conda activate desc-env

cd DESC
pip install -r devtools/dev-requirements.txt   # or just requirements.txt
pip install "jax[cuda12]" # this already install all CUDA wheels

import sys
import os
sys.path.append(os.path.abspath("/pscratch/sd/r/rgaur/DESC")) 
os.environ["LD_LIBRARY_PATH "] = ""  # before any import from desc

from desc import set_device
set_device("gpu")
from desc.equilibrium import Equilibrium

which gives

>>                           (reverse-i-search)`from desc.eq': from desc.equilibrium import Equilibrium
ERROR:2025-09-19 13:13:40,766:jax._src.xla_bridge:502: Jax plugin configuration error: Exception when calling jax_plugins.xla_cuda12.initialize()
Traceback (most recent call last):
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 201, in _version_check
    version = get_version()
              ^^^^^^^^^^^^^
RuntimeError: jaxlib/cuda/versions_helpers.cc:81: operation cusparseGetProperty(MAJOR_VERSION, &major) failed: The cuSPARSE library was not found.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax/_src/xla_bridge.py", line 500, in discover_pjrt_plugins
    plugin_module.initialize()
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 328, in initialize
    _check_cuda_versions(raise_on_first_error=True)
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 266, in _check_cuda_versions
    _version_check("cuSPARSE", cuda_versions.cusparse_get_version,
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 205, in _version_check
    raise RuntimeError(err_msg) from e
RuntimeError: Unable to load cuSPARSE. Is it installed?
WARNING:2025-09-19 13:13:40,819:jax._src.xla_bridge:872: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/pscratch/sd/r/rgaur/DESC/desc/equilibrium/__init__.py", line 3, in <module>
    from .equilibrium import EquilibriaFamily, Equilibrium
  File "/pscratch/sd/r/rgaur/DESC/desc/equilibrium/equilibrium.py", line 13, in <module>
    from desc.backend import execute_on_cpu, jnp
  File "/pscratch/sd/r/rgaur/DESC/desc/backend.py", line 31, in <module>
    if desc_config.get("kind") == "gpu" and len(jax.devices("gpu")) == 0:
                                                ^^^^^^^^^^^^^^^^^^
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax/_src/xla_bridge.py", line 1018, in devices
    return get_backend(backend).devices()
           ^^^^^^^^^^^^^^^^^^^^
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax/_src/util.py", line 298, in wrapper
    return cached(config.trace_context() if trace_context_in_key else _ignore(),
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax/_src/util.py", line 292, in cached
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax/_src/xla_bridge.py", line 952, in get_backend
    return _get_backend_uncached(platform)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax/_src/xla_bridge.py", line 933, in _get_backend_uncached
    platform = canonicalize_platform(platform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/homes/r/rgaur/.conda/envs/desc-env/lib/python3.12/site-packages/jax/_src/xla_bridge.py", line 736, in canonicalize_platform
    raise RuntimeError(f"Unknown backend: '{platform}' requested, but no "
RuntimeError: Unknown backend: 'gpu' requested, but no platforms that are instances of gpu are present. Platforms present are: cpu

So the current Perlmutter instructions don't work. What does work is this:

CONDA_OVERRIDE_CUDA="12.4" conda create --name desc-env "jax==0.6.0" "jaxlib==0.6.0=cuda12*" -c conda-forge
cd DESC
pip install --editable .

and then I don't get any error loading equilibria.

[YigitElma]

Hmm, we added the LD_LIBRARY_PATH thing exactly for this error message. Maybe the cuda version change cause the problem. I will have access to Perlmutter next week, I can try it too.

[YigitElma]

Shared this with Rahul already, for documentation purposes, sharing here again.

module load conda

conda create -n desc-env python=3.12
conda activate desc-env

cd DESC
pip install -r devtools/dev-requirements.txt   # or just requirements.txt
pip install "jax[cuda12]" # this already install all CUDA wheels

These instructions work on Perlmutter compute nodes with JAX 0.6.2 without needing LD_LIBRARY_PATH trick. The problem happens if you try to check your installation on the login node. Login node has a GPU which can be seen via nvidia-smi but for some reason, simple tests cause the error Rahul shared above (even if you use LD_LIBRARY_PATH command) (most probably due to incompatible CUDA versions installed). Just import jax; jax.devices("gpu") cause the same error. Given that this happens only on the login node and not on the compute nodes, I believe we can use the usual instructions with slight updates and document this at the end verbally.