Closed gpadiolleau closed 1 year ago
I'm speculating you have two versions of CUDNN installed. If you import torch
first, you're probably getting the version that torch
uses, and if you import jax
first, you're getting jax
's version. This might create problems if the jax
version is newer, uses more GPU memory, and you run out if you use it.
Can you try lowering XLA_PYTHON_CLIENT_MEM_FRACTION
? https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
It defaults to 0.75, try 0.7 or lower.
Hi, thanks for your help !
Yes, I have to two CUDA installations in my environment. If I run conda list cuda
and conda list cudnn
, I get the following :
# packages in environment at /home/ubuntu/anaconda3/envs/PyTorchGPU:
#
# Name Version Build Channel
cuda-cudart 11.8.89 0 nvidia
cuda-cupti 11.8.87 0 nvidia
cuda-libraries 11.8.0 0 nvidia
cuda-nvrtc 11.8.89 0 nvidia
cuda-nvtx 11.8.86 0 nvidia
cuda-runtime 11.8.0 0 nvidia
cudatoolkit 11.8.0 h6a678d5_0
pytorch-cuda 11.8 h7e8668a_5 pytorch
# packages in environment at /home/ubuntu/anaconda3/envs/PyTorchGPU:
#
# Name Version Build Channel
cudnn 8.9.2.26 cuda11_0
But, as you can see, torch
comes with its own CUDA installation pytorch-cuda
.
However, both are CUDA version 11.8 (and I can't find torch
CUDNN version)
As I said in my first message, the problem does not seem to come from import's order since importing torch
first or jax
first does not change anything : it works when I have import torch
before any jax
operations.
I do not think decreasing XL_PYTHON_CLIENT_MEM_FRACTION
will work since Memory usage
is less than total bytes in error message, but I may have misunderstood this information.
I ran the same test with XL_PYTHON_CLIENT_MEM_FRACTION=0.10
:
JAX_PLATFORM_NAME=gpu XLA_PYTHON_CLIENT_MEM_FRACTION=0.10 python -c "from jax.random import PRNGKey; print(PRNGKey(0));"
but I got the same error :
2023-07-18 11:50:19.639897: E external/xla/xla/stream_executor/cuda/cuda_dnn.cc:439] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
2023-07-18 11:50:19.639960: E external/xla/xla/stream_executor/cuda/cuda_dnn.cc:443] Memory usage: 5653528576 bytes free, 6367477760 bytes total.
Traceback (most recent call last):
File "<string>", line 1, in <module>
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/random.py", line 160, in PRNGKey
key = prng.seed_with_impl(impl, seed)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/prng.py", line 406, in seed_with_impl
return random_seed(seed, impl=impl)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/prng.py", line 690, in random_seed
return random_seed_p.bind(seeds_arr, impl=impl)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/prng.py", line 702, in random_seed_impl
base_arr = random_seed_impl_base(seeds, impl=impl)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/prng.py", line 707, in random_seed_impl_base
return seed(seeds)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/prng.py", line 936, in threefry_seed
return _threefry_seed(seed)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/pjit.py", line 250, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/pjit.py", line 163, in _python_pjit_helper
out_flat = pjit_p.bind(*args_flat, **params)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/core.py", line 2677, in bind
return self.bind_with_trace(top_trace, args, params)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/core.py", line 383, in bind_with_trace
out = trace.process_primitive(self, map(trace.full_raise, args), params)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/core.py", line 815, in process_primitive
return primitive.impl(*tracers, **params)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/pjit.py", line 1203, in _pjit_call_impl
return xc._xla.pjit(name, f, call_impl_cache_miss, [], [], donated_argnums,
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/pjit.py", line 1187, in call_impl_cache_miss
out_flat, compiled = _pjit_call_impl_python(
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/pjit.py", line 1120, in _pjit_call_impl_python
compiled = _pjit_lower(
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py", line 2323, in compile
executable = UnloadedMeshExecutable.from_hlo(
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py", line 2645, in from_hlo
xla_executable, compile_options = _cached_compilation(
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py", line 2555, in _cached_compilation
xla_executable = dispatch.compile_or_get_cached(
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/dispatch.py", line 497, in compile_or_get_cached
return backend_compile(backend, computation, compile_options,
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/dispatch.py", line 465, in backend_compile
return backend.compile(built_c, compile_options=options)
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed. Look at the errors above for more details.
The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.
--------------------
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "<string>", line 1, in <module>
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/random.py", line 160, in PRNGKey
key = prng.seed_with_impl(impl, seed)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/prng.py", line 406, in seed_with_impl
return random_seed(seed, impl=impl)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/prng.py", line 690, in random_seed
return random_seed_p.bind(seeds_arr, impl=impl)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/core.py", line 380, in bind
return self.bind_with_trace(find_top_trace(args), args, params)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/core.py", line 383, in bind_with_trace
out = trace.process_primitive(self, map(trace.full_raise, args), params)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/core.py", line 815, in process_primitive
return primitive.impl(*tracers, **params)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/prng.py", line 702, in random_seed_impl
base_arr = random_seed_impl_base(seeds, impl=impl)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/prng.py", line 707, in random_seed_impl_base
return seed(seeds)
File "/home/ubuntu/anaconda3/envs/PyTorchGPU/lib/python3.8/site-packages/jax/_src/prng.py", line 936, in threefry_seed
return _threefry_seed(seed)
jaxlib.xla_extension.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed. Look at the errors above for more details.
Is it possible that the problem is caused by my jaxlib
installation ?
I have jaxlib
version 0.4.13+cuda11.cudnn86
since this is the latest version for CUDA 11, but my environment run CUDNN 8.9 for CUDA 11.
I installed jaxlib
version 0.4.13+cuda11.cudnn86
since version for CUDNN 8.9 are only available for CUDA 12 according to https://storage.googleapis.com/jax-releases/jax_cuda_releases.html, but torch
is only available for CUDA 11 (hahahaha....)
p.s. : Is it normal that decreasing XLA_PYTHON_CLIENT_MEM_FRACTION
results in increasing Memory usage
in error message ?
I suspect your multiple CUDA installations are confusing something. Note that torch
and jax
actually have conflicting CUDA requirements.
I recommend using separate virtual environments for torch
and jax
, and installing jax
using the cuda12_pip
wheels per the README.
I'm not sure there's any action we can take here.
I thought that installing jax
and torch
for the same CUDA and CUDNN version would have work.
I created a new environment and install jax
with pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
.
Re-test JAX_PLATFORM_NAME=gpu python -c "from jax.random import PRNGKey; print(PRNGKey(0));"
and it works like a charm !
I'll try later to get an environment with local CUDA and CUDNN with torch
and jax
.
Thanks for your help, and btw thanks for making DL libraries with Nvidia GPU supports simpler to install ! I remember my passed long hours battles against Tensorflow and Nvidia drivers ! haha
FWIW I've bumped into the same error, whilst needing both PyTorch and JAX in the same environment. (I'm using PyTorch just because JAX doesn't provide any dataloaders.)
The fix seems to be running the following before the rest of my program:
import jax
jax.random.PRNGKey(0) # Initialise JAX's backend by performing a nontrivial JAX operation.
import torch # Afterwards!
Notably this is not the same fix described above, so I wanted to share it here in case it helped anyone out.
@patrick-kidger Yes, that's likely because JAX wants a newer CUDA/CUDNN release than PyTorch does, and by loading the newer one first, things probably work.
The safer option is to install a CPU-only pytorch, though.
There's not a ton we can do about this, other than perhaps adding some better errors if you get it wrong.
I strongly suspect that https://github.com/google/jax/commit/9404518201c5ac8af6c85ecdf12a7cc34c102585 will help catch this case. We will now raise an exception if we detect that the versions of the CUDA libraries loaded are too old.
Description
I have a conda environnement named PyTorchGPU where I installed torchGPU at first and then installed JAX and jaxlib with Nvidia GPU support. It seems running JAX only does not allow for cudnn to use the GPU but if I import torch first, it magically works !
Running JAX only
When I run a simple test on GPU in the terminal :
(PyTorchGPU) ubuntu@ubuntu-GL62MVR-7RFX:~$ JAX_PLATFORM_NAME=gpu python -c "from jax.random import PRNGKey; PRNGKey(10)"
, I get the following error :However running the following test
(PyTorchGPU) ubuntu@ubuntu-GL62MVR-7RFX:~$ JAX_PLATFORM_NAME=gpu python -c "import jax;print(f'JAX available devices :',jax.devices())"
gives me the following result asserting JAX detects my GPU :JAX available devices : [gpu(id=0)]
First
import torch
and run JAXRunning the following test on GPU in terminal :
(PyTorchGPU) ubuntu@ubuntu-GL62MVR-7RFX:~$ JAX_PLATFORM_NAME=gpu python -c "import torch;import jax;print(jax.random.PRNGKey(10))"
gives me the expected result :[ 0 10]
It is not related imports order since running
JAX_PLATFORM_NAME=gpu python -c "import jax;import torch;print(jax.random.PRNGKey(10))"
gives the same result.Any ideas ?
It really seems to be a problem with cudnn kind of "activation" through torch. Maybe 'cause torch comes with it cudnn installation ? Here is my
conda list cuda
result :Does
pytorch-cuda
overwrite my local codatoolkit installation ? And then I need toimport torch
first to runpytorch-cuda
as backend for JAX ?Thanks for any help, I need to understand this ! :)
What jax/jaxlib version are you using?
jax 0.4.13 / jaxlib 0.4.13+cuda11.cudnn86
Which accelerator(s) are you using?
GPU
Additional system info
Python 3.8.17, Ubuntu 22.04.2 LTS
NVIDIA GPU info