Incompatible CUDA versions installed from Pip

[rosario-purple]

Description

When using the default install command:

pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

it tries to install CUDA 12.3 libraries. However, the most recent stable version of PyTorch (2.1.2) is pinned to 12.1, and that's what I have on my machine. (I'm compiling flash-attn, TransformerEngine and MS-AMP on this machine after installing the main package dependencies, and the compilation appears to be sensitive to CUDA version.)

What jax/jaxlib version are you using?

0.4.23

Which accelerator(s) are you using?

GPU, Nvidia A100

Additional system info?

1.24.4 3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0] uname_result(system='Linux', node='7e72bd4e-01', release='5.15.0-91-generic', version='#101-Ubuntu SMP Tue Nov 14 13:30:08 UTC 2023', machine='x86_64')

NVIDIA GPU info

Mon Jan 22 11:32:43 2024
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 545.23.08 Driver Version: 545.23.08 CUDA Version: 12.3 |
|-----------------------------------------+----------------------+----------------------+
| 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 A100-SXM4-80GB On | 00000000:11:00.0 Off | 0 |
| N/A 29C P0 62W / 400W | 5MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
| 1 NVIDIA A100-SXM4-80GB On | 00000000:12:00.0 Off | 0 |
| N/A 27C P0 60W / 400W | 5MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
| 2 NVIDIA A100-SXM4-80GB On | 00000000:13:00.0 Off | 0 |
| N/A 26C P0 61W / 400W | 5MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
| 3 NVIDIA A100-SXM4-80GB On | 00000000:14:00.0 Off | 0 |
| N/A 28C P0 62W / 400W | 5MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
| 4 NVIDIA A100-SXM4-80GB On | 00000000:21:00.0 Off | 0 |
| N/A 29C P0 61W / 400W | 5MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
| 5 NVIDIA A100-SXM4-80GB On | 00000000:22:00.0 Off | 0 |
| N/A 27C P0 61W / 400W | 5MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
| 6 NVIDIA A100-SXM4-80GB On | 00000000:23:00.0 Off | 0 |
| N/A 26C P0 62W / 400W | 5MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+
| 7 NVIDIA A100-SXM4-80GB On | 00000000:24:00.0 Off | 0 |
| N/A 29C P0 62W / 400W | 5MiB / 81920MiB | 0% Default |
| | | Disabled |
+-----------------------------------------+----------------------+----------------------+

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

[jakevdp]

Hi - thanks for the question. You'll find some relevant info here: #18032 (comment)

Our general approach is to provide two JAX CUDA builds: one at the most recent release (currently 12.3) and one older version that seeks to maintain compatibility with pytorch's requirements. Currently that is our cuda11 build.

[jakevdp]

I'm going to close this, since it's essentially a duplicate of #18032. Thanks!

[hawkinsp]

I'll note that the safest thing to do in general is use two separate venvs. Different frameworks have different version requirements. JAX tracks CUDA versions faster than PyTorch does.