JAX errors with GPU options

[dal281726-collab]

[Edit] I was trying to run on an H100 GPU node.
I am a rookie to Dolphin. I installed Dolphin and tried to follow the tutorial. It's pretty straightforward to install and run it on CPU nodes, so I wonder if it's possible to use GPUs to accelerate it. I followed the instructions from Dolphin tutorials, then I first tried submitting a job, and it looked like the job wasn't running on the GPU so tried to see if the GPU was visible to dolphin/JAX.

For dolphin, I ran (on a GPU node)

$ python3
>>> import dolphin
>>> dolphin.show_versions()

and it printed out gpu_is_available: True, so it looked good. Then I tried

$ python3
>>> import jax
>>> print(jax.devices())

It printed out WARNING:2025-08-13 11:49:16,983: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. [CpuDevice(id=0)]
So the GPU was not seen by JAX.

After that, I created a new environment and installed JAX from the original repo. I tried running print(jax.devices()) again and JAX actually saw the GPU with no problem:

$ python3 -c "import jax; print(jax.devices())"
[CudaDevice(id=0)]

I thought it should be okay now. Then I tried installing dolphin again, and ran python3 -c "import jax; print(jax.devices())" again, but it still threw the same error WARNING:2025-08-13 12:45:19,983: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. [CpuDevice(id=0)]

No matter how I tried uninstall/reinstall/upgrade/downgrade JAX several times, it didn't work. It only works with JAX alone. If you have any suggestions please let me know!

[scottstanie]

Interesting... can you run python -c 'import dolphin; dolphin.show_versions()'?
How did you install dolphin/jax?

I'm wondering if there's some conda issue with jax+jaxlib on certain GPUs.

[dal281726-collab]

I ran that command without any issues. I installed dolphin by mamba install -c conda-forge dolphin, and apparently jax was installed as a dependency so I didn't have to install it. I also double checked:

$ pip show jax jaxlib
Name: jax
Version: 0.6.2
Summary: Differentiate, compile, and transform Numpy code.
Home-page: https://github.com/jax-ml/jax
Author: JAX team
Author-email: [email protected]
License: Apache-2.0
Location: [redacted]
Requires: jaxlib, ml_dtypes, numpy, opt_einsum, scipy
Required-by: dolphin
---
Name: jaxlib
Version: 0.6.2
Summary: XLA library for JAX
Home-page: https://github.com/jax-ml/jax
Author: JAX team
Author-email: [email protected]
License: Apache-2.0
Location: [redacted]
Requires: ml_dtypes, numpy, scipy
Required-by: jax

jaxlib is always there.

[scottstanie]

I see... maybe you could try installing numba, cudatoolkit, and Jax into an environment and see if that works? If you're able to do that, you can just add dolphin in after. Sorry that there's not a very simple way while we still have a few numba things in there, it can be annoying via conda forge

[dal281726-collab]

Thank you very much for the suggestion. Now I ran into a different issue: My GPU got allocated some processes but the GPU was idle, it wasn't being utilized at all

[scottstanie]

oh awesome!

so now there's a different, better problem haha.

Since I've never run it on an H100, i've got a few questions:

• did it speed things up at all compared to the CPU (if you've never tried CPU only, disregard)
• were you processing multiple OPERA CSLC burst ids? or one big stack, like what comes out of isce2?

From my own usage on an A100, it's true that each separate wrapped_phase process does not really come close to 100% utilization, unless you have set --phase-linking-options.ministack-size to be quite large, and/or have set the --worker-setting.block-size large. I'll usually set the environment variable export XLA_PYTHON_CLIENT_MEM_FRACTION=0.15 or so, and run 3-5 separate bursts at a time to help utilization.

The phase linking process isn't super GPU intensive compared to, e.g., a large language model, so it seems like the newer GPUs are way over powered for a little phase linking.

I still see at least 3-5x speeds for phase linking, and 5-10x speedups for the "timeseries inversion" step, especially on a large stack. but i have suspect it would be very hard to get close to 100% utilization

[scottstanie]

I also don't want to get your hopes up about dolphin's GPU ability, since the phase unwrapping step (often the slowest part of the full workflow) does not use it. I have yet to see a real GPU-compatible phase unwrapping algorithm. The only one close that I know of it my small experiment here, which would need to have the numpy matrix multiplies get converted to JAX

[scottstanie]

you could try installing dolphin as you did, and then do a pip install jax jaxlib to see if it better detect your GPU environment.
But it also might be some bad interaction between numba (which gets cudatoolkit) and jax...

[dal281726-collab]

I also tried that approach! And pretty much I got the same errors. Standalone JAX ran perfectly though, maybe just the binding is the problem .