XLA error when running on GPU #14586

chufanchen · 2023-02-20T11:13:45Z

chufanchen
Feb 20, 2023

Description

On version 0.4.5 I receive the following error when I run

>>> from jax import random
>>> import jax.numpy as jnp
>>> key = random.PRNGKey(0)

2023-02-20 18:59:00.324153: W external/org_tensorflow/tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:397] There was an error before creating cudnn handle: cudaGetErrorName symbol not found. : cudaGetErrorString symbol not found.
2023-02-20 18:59:01.434511: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:429] Could not create cudnn handle: CUDNN_STATUS_NOT_INITIALIZED
2023-02-20 18:59:01.434629: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:438] Possibly insufficient driver version: 525.85.12
2023-02-20 18:59:01.435448: E external/org_tensorflow/tensorflow/compiler/xla/status_macros.cc:57] INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:627) dnn != nullptr
*** Begin stack trace ***
	_PyObject_MakeTpCall

	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	PyObject_Call
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault

	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall

	PyObject_Call
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall


	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	PyObject_Call
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault

	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault

	PyEval_EvalCode



	_PyRun_InteractiveLoopObject

	PyRun_AnyFileExFlags

	Py_BytesMain
	__libc_start_main

*** End stack trace ***

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/lab304/jax/jax/_src/random.py", line 136, in PRNGKey
    key = prng.seed_with_impl(impl, seed)
  File "/home/lab304/jax/jax/_src/prng.py", line 267, in seed_with_impl
    return random_seed(seed, impl=impl)
  File "/home/lab304/jax/jax/_src/prng.py", line 570, in random_seed
    return random_seed_p.bind(seeds_arr, impl=impl)
  File "/home/lab304/jax/jax/_src/core.py", line 342, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
  File "/home/lab304/jax/jax/_src/core.py", line 345, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/home/lab304/jax/jax/_src/core.py", line 788, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/home/lab304/jax/jax/_src/prng.py", line 582, in random_seed_impl
    base_arr = random_seed_impl_base(seeds, impl=impl)
  File "/home/lab304/jax/jax/_src/prng.py", line 587, in random_seed_impl_base
    return seed(seeds)
  File "/home/lab304/jax/jax/_src/prng.py", line 822, in threefry_seed
    lax.shift_right_logical(seed, lax_internal._const(seed, 32)))
  File "/home/lab304/jax/jax/_src/lax/lax.py", line 512, in shift_right_logical
    return shift_right_logical_p.bind(x, y)
  File "/home/lab304/jax/jax/_src/core.py", line 342, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
  File "/home/lab304/jax/jax/_src/core.py", line 345, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/home/lab304/jax/jax/_src/core.py", line 788, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/home/lab304/jax/jax/_src/dispatch.py", line 124, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args),
  File "/home/lab304/jax/jax/_src/util.py", line 253, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
  File "/home/lab304/jax/jax/_src/util.py", line 246, in cached
    return f(*args, **kwargs)
  File "/home/lab304/jax/jax/_src/dispatch.py", line 203, in xla_primitive_callable
    compiled = _xla_callable_uncached(lu.wrap_init(prim_fun), device, None,
  File "/home/lab304/jax/jax/_src/dispatch.py", line 356, in _xla_callable_uncached
    return computation.compile(_allow_propagation_to_outputs=allow_prop).unsafe_call
  File "/home/lab304/jax/jax/_src/interpreters/pxla.py", line 3281, in compile
    executable = self._compile_unloaded(
  File "/home/lab304/jax/jax/_src/interpreters/pxla.py", line 3252, in _compile_unloaded
    return UnloadedMeshExecutable.from_hlo(
  File "/home/lab304/jax/jax/_src/interpreters/pxla.py", line 3539, in from_hlo
    xla_executable = dispatch.compile_or_get_cached(
  File "/home/lab304/jax/jax/_src/dispatch.py", line 1095, in compile_or_get_cached
    return backend_compile(backend, serialized_computation, compile_options,
  File "/home/lab304/jax/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/lab304/jax/jax/_src/dispatch.py", line 1040, in backend_compile
    return backend.compile(built_c, compile_options=options)
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:627) dnn != nullptr

Ubuntu 20.04, CUDA 12.0, CuDNN 8.8
GPU: 4 RTX 2080ti

What jax/jaxlib version are you using?

jax 0.4.5/jaxlib 0.4.5

I install wax from source

(jax-exp) ➜  jax git:(main) python build/build.py --enable_cuda --bazel_options "--jobs 56"

     _   _  __  __
    | | / \ \ \/ /
 _  | |/ _ \ \  /
| |_| / ___ \/  \
 \___/_/   \/_/\_\


Bazel binary path: ./bazel-5.1.1-linux-x86_64
Bazel version: 5.1.1
Python binary path: /home/lab304/.conda/envs/jax-exp/bin/python
Python version: 3.10
NumPy version: 1.24.2
MKL-DNN enabled: yes
Target CPU: x86_64
Target CPU features: release
CUDA enabled: yes
NCCL enabled: yes
TPU enabled: no
Remote TPU enabled: no
ROCm enabled: no
Plugin device enabled: no

Which accelerator(s) are you using?

GPU

Additional system info

Python 3.10

NVIDIA GPU info

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| 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 ...  On   | 00000000:02:00.0 Off |                  N/A |
| 27%   27C    P8     1W / 250W |   8680MiB / 11264MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce ...  On   | 00000000:03:00.0 Off |                  N/A |
| 27%   27C    P8     1W / 250W |    400MiB / 11264MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  NVIDIA GeForce ...  On   | 00000000:82:00.0 Off |                  N/A |
| 27%   26C    P8    11W / 250W |    400MiB / 11264MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   3  NVIDIA GeForce ...  On   | 00000000:83:00.0 Off |                  N/A |
| 27%   28C    P8    17W / 250W |    400MiB / 11264MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1800      G   /usr/lib/xorg/Xorg                  9MiB |
|    0   N/A  N/A      2063      G   /usr/bin/gnome-shell                3MiB |
|    0   N/A  N/A     22551      C   ...a/envs/jax-exp/bin/python     8662MiB |
|    1   N/A  N/A      1800      G   /usr/lib/xorg/Xorg                  4MiB |
|    1   N/A  N/A     22551      C   ...a/envs/jax-exp/bin/python      392MiB |
|    2   N/A  N/A      1800      G   /usr/lib/xorg/Xorg                  4MiB |
|    2   N/A  N/A     22551      C   ...a/envs/jax-exp/bin/python      392MiB |
|    3   N/A  N/A      1800      G   /usr/lib/xorg/Xorg                  4MiB |
|    3   N/A  N/A     22551      C   ...a/envs/jax-exp/bin/python      392MiB |
+-----------------------------------------------------------------------------+

Answered by chufanchen

Feb 21, 2023

I downgraded the cuda and reinstalled cudnn. After that, it works.

View full answer

chufanchen · 2023-02-21T07:00:19Z

chufanchen
Feb 21, 2023
Author

I downgraded the cuda and reinstalled cudnn. After that, it works.

2 replies

jakevdp Feb 21, 2023
Maintainer

This seems like the right solution – as of jaxlib v0.4.5, I don't believe that the pre-built wheels are compatible with CUDA 12.0.

nouiz Mar 16, 2023
Collaborator

Note, we make the error message shorter: tensorflow/tensorflow#59915
So hopefully, it will be easier to understand.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

XLA error when running on GPU #14586

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Replies: 1 comment 2 replies

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Select a reply

Uh oh!

XLA error when running on GPU #14586

Uh oh!

Uh oh!

chufanchen Feb 20, 2023

Description

Replies: 1 comment · 2 replies

Uh oh!

Uh oh!

chufanchen Feb 21, 2023 Author

Uh oh!

jakevdp Feb 21, 2023 Maintainer

Uh oh!

nouiz Mar 16, 2023 Collaborator

chufanchen
Feb 20, 2023

Replies: 1 comment 2 replies

chufanchen
Feb 21, 2023
Author

jakevdp Feb 21, 2023
Maintainer

nouiz Mar 16, 2023
Collaborator