Jax crashes after upgrade #15260

terafo · 2023-03-28T12:17:35Z

Description

All operations just outright crash after upgrade. For example

import jax
jax.numpy.zeros((5,5))

Outputs this error:

2023-03-28 15:01:11.001453: W external/xla/xla/stream_executor/cuda/cuda_dnn.cc:397] There was an error before creating cudnn handle: cudaErrorNotSupported : operation not supported
Array([0., 0., 0., 0., 0.], dtype=float32)

This, at least outputs array(fallback path, I presume). If I try

import jax
jax.random.normal(jax.random.PRNGKey(0),(5,5))

I get this error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 137, in PRNGKey
    key = prng.seed_with_impl(impl, seed)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 270, in seed_with_impl
    return random_seed(seed, impl=impl)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 561, in random_seed
    return random_seed_p.bind(seeds_arr, impl=impl)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 573, in random_seed_impl
    base_arr = random_seed_impl_base(seeds, impl=impl)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 578, in random_seed_impl_base
    return seed(seeds)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 817, in threefry_seed
    k2 = convert(jnp.bitwise_and(seed, np.uint32(0xFFFFFFFF)))
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/pjit.py", line 248, in cache_miss
    outs, out_flat, out_tree, args_flat = _python_pjit_helper(
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/pjit.py", line 195, in _python_pjit_helper
    out_flat = pjit_p.bind(*args_flat, **params)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 2591, in bind
    return self.bind_with_trace(top_trace, args, params)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 362, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 816, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/pjit.py", line 1269, in _pjit_call_impl
    return compiled.unsafe_call(*args)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 1934, in __call__
    results = self.xla_executable.execute_sharded(input_bufs)
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: INVALID_ARGUMENT: stream is uninitialized or in an error state

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 "<stdin>", line 1, in <module>
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 137, in PRNGKey
    key = prng.seed_with_impl(impl, seed)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 270, in seed_with_impl
    return random_seed(seed, impl=impl)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 561, in random_seed
    return random_seed_p.bind(seeds_arr, impl=impl)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 359, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 362, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 816, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 573, in random_seed_impl
    base_arr = random_seed_impl_base(seeds, impl=impl)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 578, in random_seed_impl_base
    return seed(seeds)
  File "/home/terafo/jt/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 817, in threefry_seed
    k2 = convert(jnp.bitwise_and(seed, np.uint32(0xFFFFFFFF)))
jaxlib.xla_extension.XlaRuntimeError: INVALID_ARGUMENT: stream is uninitialized or in an error state

I've got this issue when I tried to use both local and pip cuda.

What jax/jaxlib version are you using?

0.4.7

Which accelerator(s) are you using?

GPU

Additional system info

Python 3.10 in WSL with CUDNN 8.8.1 and CUDA 12.1

NVIDIA GPU info

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 530.41.03              Driver Version: 531.41       CUDA Version: 12.1     |
|-----------------------------------------+----------------------+----------------------+
| 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 GTX 1080 Ti      On | 00000000:09:00.0  On |                  N/A |
| 20%   46C    P8               13W / 250W|  10019MiB / 11264MiB |      2%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A        22      G   /Xwayland                                 N/A      |
|    0   N/A  N/A       487      C   /python3.10                               N/A      |
+---------------------------------------------------------------------------------------+

The text was updated successfully, but these errors were encountered:

hawkinsp · 2023-03-28T12:25:06Z

We don't ourselves support WSL, but I expect it should work. Do any CUDA 12 apps work in this setup?

terafo · 2023-03-28T12:44:50Z

I successfully compiled and ran one cuda sample(transpose) and one cudnn sample(mnistCUDNN) with no issues. As far as I can tell torch runs fine as well. When I used CUDA 11.8+cudnn8.6 Jax worked just fine.

terafo · 2023-03-28T18:46:08Z

I managed to get it working with CUDA 11, so problem is definitely localized to CUDA 12. Maybe it has something to do with me using Pascal GPU?

hawkinsp · 2023-03-29T13:54:41Z

Unfortunately we don't ourselves test or support Windows or WSL2, but if you can figure out the problem we certainly would accept contributions to fix it!

mikelaud · 2023-03-29T19:21:59Z

Hi,
just tested jax + wsl2 and found no problem:

WSL2 + Ubuntu-20.04 + jax-0.4.7

https://github.com/google/jax#installation
wsl --install -d Ubuntu-20.04
sudo su

apt update
apt -y full-upgrade

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

python3

import jax
jax.numpy.zeros((5,5))
Array([[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.]], dtype=float32)
jax.default_backend()
'gpu'

Window Edition: Windows 11 Pro
Version: 21H2
OS build: 22000.1696

Studio Driver
Version 531.41
NVIDIA GeForce RTX 3090

WSL2 + Ubuntu-22.04 + jax-0.4.7

https://github.com/google/jax#installation
wsl --install -d Ubuntu-22.04
sudo su

apt update
apt -y full-upgrade

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

python3

import jax
jax.numpy.zeros((5,5))
Array([[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.]], dtype=float32)
jax.default_backend()
'gpu'

Window Edition: Windows 11 Pro
Version: 21H2
OS build: 22000.1696

Studio Driver
Version 531.41
NVIDIA GeForce RTX 3090

terafo · 2023-03-29T20:03:48Z

@mikelaud just tried same sequence with fresh install of Ubuntu 20.04. I receive the same error. Not sure what's the issue. I have 531.41 driver(though game ready), Windows 11 Pro 22H2 with build version being 22621.1413. And my GPU is 1080 Ti.
Update: switched to studio driver, still same error. I feel cursed.

terafo · 2023-03-30T20:14:09Z

Ok, now I know that jax is broken on Pascal with WSL with any version of CUDA(less features are broken with CUDA 11, more with CUDA 12). I tested a couple of commits and figured that after XLA got upgraded from 0f31407ee498e6dba242d03f8d382ebcfcc61790 to 79ca8d03c296ede04dc9a86ce9dde79ed909dda8 all my issues started. Not sure which version of XLA is to blame, not sure how to build jax against different versions of XLA.

mikelaud · 2023-03-31T09:34:24Z

@terafo Well, it looks like one should create a bug within the XLA framework. After the fix on the XLA side, it will be possible to raise the version for JAX.

terafo · 2023-05-10T09:32:37Z

With recent 0.4.9 upgrade the issue is resolved since it uses more recent XLA version, in which this issue is fixed.

terafo added the bug Something isn't working label Mar 28, 2023

google-ml-butler bot added the type:support label Mar 28, 2023

hawkinsp added the NVIDIA GPU Issues specific to NVIDIA GPUs label Mar 28, 2023

hawkinsp added contributions welcome The JAX team has not prioritized work on this. Community contributions are welcome. Windows Issues related to JAX on Microsoft Windows labels Mar 29, 2023

terafo mentioned this issue Mar 31, 2023

JAX built against XLA crashes on Pascal GPUs in WSL after commit 16e953a921c150dcaf012299a52d771e4bb425c2 openxla/xla#2249

Closed

terafo closed this as completed May 10, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Jax crashes after upgrade #15260

Jax crashes after upgrade #15260

terafo commented Mar 28, 2023 •

edited

Loading

hawkinsp commented Mar 28, 2023

terafo commented Mar 28, 2023 •

edited

Loading

terafo commented Mar 28, 2023

hawkinsp commented Mar 29, 2023

mikelaud commented Mar 29, 2023 •

edited

Loading

terafo commented Mar 29, 2023 •

edited

Loading

terafo commented Mar 30, 2023

mikelaud commented Mar 31, 2023

terafo commented May 10, 2023

Jax crashes after upgrade #15260

Jax crashes after upgrade #15260

Comments

terafo commented Mar 28, 2023 • edited Loading

Description

What jax/jaxlib version are you using?

Which accelerator(s) are you using?

Additional system info

NVIDIA GPU info

hawkinsp commented Mar 28, 2023

terafo commented Mar 28, 2023 • edited Loading

terafo commented Mar 28, 2023

hawkinsp commented Mar 29, 2023

mikelaud commented Mar 29, 2023 • edited Loading

terafo commented Mar 29, 2023 • edited Loading

terafo commented Mar 30, 2023

mikelaud commented Mar 31, 2023

terafo commented May 10, 2023

terafo commented Mar 28, 2023 •

edited

Loading

terafo commented Mar 28, 2023 •

edited

Loading

mikelaud commented Mar 29, 2023 •

edited

Loading

terafo commented Mar 29, 2023 •

edited

Loading