CUDA_ERROR_INVALID_IMAGE: device kernel image is invalid

[jacktang]

Hello,

I try to install JAX on my PC using the commend

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

conda list | grep jax
# packages in environment at /home/dev/miniconda3/envs/jax:
jax                       0.4.13                   pypi_0    pypi
jaxlib                    0.4.13+cuda12.cudnn89          pypi_0    pypi

The PC OS is Ubuntu 20.04 and the graphic card information:

nvidia-smi
Fri Jul 21 10:56:53 2023
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.54.03              Driver Version: 535.54.03    CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| 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 RTX 3090 Ti     On  | 00000000:01:00.0 Off |                  Off |
|  0%   46C    P8              27W / 450W |  18470MiB / 24564MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      3026      C   .../dev/miniconda3/envs/jax/bin/python    18462MiB |
+---------------------------------------------------------------------------------------+

And I run the test code:

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random

key = random.PRNGKey(0)
x = random.normal(key, (10,))
print(x)

some errors come up:

2023-07-21 11:03:18.706602: E external/xla/xla/stream_executor/gpu/asm_compiler.cc:114] *** WARNING *** You are using ptxas 10.1.243, which is older than 11.1. ptxas before 11.1 is known to miscompile XLA code, leading to incorrect results or invalid-address errors.

2023-07-21 11:03:18.707395: W external/xla/xla/stream_executor/gpu/asm_compiler.cc:231] Falling back to the CUDA driver for PTX compilation; ptxas does not support CC 8.6
2023-07-21 11:03:18.707401: W external/xla/xla/stream_executor/gpu/asm_compiler.cc:234] Used ptxas at ptxas
2023-07-21 11:03:18.713648: E external/xla/xla/stream_executor/cuda/cuda_driver.cc:548] failed to load PTX text as a module: CUDA_ERROR_INVALID_IMAGE: device kernel image is invalid
2023-07-21 11:03:18.713664: E external/xla/xla/stream_executor/cuda/cuda_driver.cc:553] error log buffer (63 bytes): error   : Binary format for key='0', ident='' is not recognize
2023-07-21 11:03:18.713708: W external/xla/xla/service/gpu/runtime/support.cc:58] Intercepted XLA runtime error:
INTERNAL: Failed to load PTX text as a module: CUDA_ERROR_INVALID_IMAGE: device kernel image is invalid
2023-07-21 11:03:18.713752: E external/xla/xla/pjrt/pjrt_stream_executor_client.cc:2461] Execution of replica 0 failed: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.func.launch' failed: Failed to load PTX text as a module: CUDA_ERROR_INVALID_IMAGE: device kernel image is invalid; current tracing scope: fusion; current profiling annotation: XlaModule:#hlo_module=jit__threefry_seed,program_id=0#.

I've searched the problem using google and still can't solved it. Could you please help me ? Thanks!

[jacktang]

I upgraded ptxas by using "conda install cuda -c nvidia", and the errors list below:

2023-07-21 11:34:18.055947: E external/xla/xla/stream_executor/cuda/cuda_driver.cc:548] failed to load PTX text as a module: CUDA_ERROR_INVALID_IMAGE: device kernel image is invalid
2023-07-21 11:34:18.056010: E external/xla/xla/stream_executor/cuda/cuda_driver.cc:553] error log buffer (63 bytes): error   : Binary format for key='0', ident='' is not recognize
2023-07-21 11:34:18.056118: W external/xla/xla/service/gpu/runtime/support.cc:58] Intercepted XLA runtime error:
INTERNAL: Failed to load PTX text as a module: CUDA_ERROR_INVALID_IMAGE: device kernel image is invalid
2023-07-21 11:34:18.056196: E external/xla/xla/pjrt/pjrt_stream_executor_client.cc:2461] Execution of replica 0 failed: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.func.launch' failed: Failed to load PTX text as a module: CUDA_ERROR_INVALID_IMAGE: device kernel image is invalid; current tracing scope: fusion; current profiling annotation: XlaModule:#hlo_module=jit__threefry_seed,program_id=0#.

[jacktang]

well, I reinstall jax and jaxlib, errors are fixed.

[kwsp]

I'm having the same problem - how did you reinstall? I tried reinstalling a few times and am still getting the same error.

[jacktang]

I reinstalled cuda apt package, python packages and jax packages and fixed the issue