Gilad Sivan

Thanks for the update, I'll check it right away

@hawkinsp After executing `pip install jax jaxlib jax-cuda12-plugin jax-cuda12-pjrt ` I'm trying to run the following code: ``` import os os.environ["CUDA_VISIBLE_DEVICES"] = "0" # Specify the index of the GPU...

@nouiz does this mean I can't yet run jax with GPU acceleration on GH200?

Last week I used the docker jax:jax here: https://github.com/NVIDIA/JAX-Toolbox I don't mind trying it again. or do u mean nvcr.io/nvidia/jax:24.04-maxtext-py3 (from here: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/jax)?

> Last week I used the docker jax:jax here: https://github.com/NVIDIA/JAX-Toolbox I don't mind trying it again. or do u mean nvcr.io/nvidia/jax:24.04-maxtext-py3 (from here: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/jax)? @nouiz

``` $ docker pull nvcr.io/nvidia/jax:24.04-maxtext-py3 24.04-maxtext-py3: Pulling from nvidia/jax no matching manifest for linux/arm64/v8 in the manifest list entries ``` @nouiz

``` gilad@gracehopper:~$ docker pull ghcr.io/nvidia/jax:jax jax: Pulling from nvidia/jax ``` works @nouiz

> ``` > gilad@gracehopper:~$ docker pull ghcr.io/nvidia/jax:jax > jax: Pulling from nvidia/jax > ``` > > works @nouiz gilad@gracehopper:~$ docker run -it --gpus all ghcr.io/nvidia/jax:jax ========== == CUDA == ==========...

@nouiz I still get: root@f85914843395:~# python jax_program.py 2024-05-21 18:06:49.257055: E external/xla/xla/stream_executor/cuda/cuda_driver.cc:282] failed call to cuInit: CUDA_ERROR_SYSTEM_NOT_READY: system not yet initialized Traceback (most recent call last): File "/usr/local/lib/python3.10/dist-packages/jax/_src/xla_bridge.py", line 679, in...

@nouiz u we're right! My colleague fixed the issue: ``` (base) nikola@gracehopper:~$ sudo nvidia-smi mig -lgi +-------------------------------------------------------+ | GPU instances: | | GPU Name Profile Instance Placement | | ID...