FIx for "Couldn't invoke ptxas --version" with cuda-11.3 and jaxlib 0.1.66+cuda111

[hrbigelow]

Hi all,

Just wanted to share my solution to the "Couldn't invoke ptxas --version" error that I got after a recent install of jax using cuda-11.3. TLDR, I needed to install nvidia-cuda-toolkit, which provides /usr/bin/ptxas. The often mentioned export XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/local/cuda-11.3 didn't work.

Details:

With my environment of cuda-11.3 and jaxlib 0.1.66+cuda111, I launched:

import jax
key = jax.random.PRNGkey(42)

and got an error Couldn't invoke ptxas --version. (details below)

Confusingly, I had:

/usr/share/cuda-11.0/bin/ptxas

but no bin directory under /usr/share/cuda-11.3.

Setting: export XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/local/cuda-11.0

did not help, but doing:

cd /usr/share/cuda-11.3 && ln -s /usr/share/cuda-11.0/bin bin

did work. However, I also found that another package, cuda-nvcc-11-0 provides ptxas, but only for cuda-11.0. I hadn't had it installed for 11.3:

(trax) henry@henry-gs65:ai$ dpkg -S /usr/local/cuda-11.0/bin/ptxas
cuda-nvcc-11-0: /usr/local/cuda-11.0/bin/ptxas

Ultimately, I did:

sudo apt install nvidia-cuda-toolkit

which provides /usr/bin/ptxas.

Once nvidia-cuda-toolkit is installed, there is no need to do either the:

export XLA_FLAGS=...

or the cd /usr/share/cuda-11.3 && ln -s /usr/share/cuda-11.0/bin bin

(trax) henry@henry-gs65:ai$ pip freeze | grep jax
jax==0.2.13
jaxlib==0.1.66+cuda111
(trax) henry@henry-gs65:ai$ nvidia-smi
Wed May 26 19:56:09 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 465.19.01    Driver Version: 465.19.01    CUDA Version: 11.3     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |

>>> os.environ.get('XLA_FLAGS')
'--xla_gpu_cuda_data_dir=/usr/local/cuda-11.0'
>>> import jax
>>> key = jax.random.PRNGKey(42)
2021-05-26 19:21:36.486473: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:169] XLA service 0x56504e05ccf0 initialized for platform Interpreter (this does not guarantee that XLA will be used). Devices:
2021-05-26 19:21:36.486542: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (0): Interpreter, <undefined>
2021-05-26 19:21:36.509052: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:169] XLA service 0x56504e06b4c0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2021-05-26 19:21:36.509084: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (0): Host, Default Version
2021-05-26 19:21:36.568099: I external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-05-26 19:21:36.568510: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:169] XLA service 0x56504df3a750 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2021-05-26 19:21:36.568526: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:177]   StreamExecutor device (0): NVIDIA GeForce GTX 1070 with Max-Q Design, Compute Capability 6.1
2021-05-26 19:21:36.568768: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/gpu_device.cc:298] Using BFC allocator.
2021-05-26 19:21:36.568804: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/gpu_device.cc:257] XLA backend allocating 7210716364 bytes on device 0 for BFCAllocator.
2021-05-26 19:21:36.569173: I external/org_tensorflow/tensorflow/stream_executor/tpu/tpu_platform_interface.cc:74] No TPU platform found.
2021-05-26 19:21:36.654318: E external/org_tensorflow/tensorflow/core/platform/default/subprocess.cc:304] Start cannot spawn child process: No such file or directory
2021-05-26 19:21:36.654343: W external/org_tensorflow/tensorflow/stream_executor/gpu/asm_compiler.cc:56] Couldn't invoke ptxas --version
2021-05-26 19:21:36.654619: E external/org_tensorflow/tensorflow/core/platform/default/subprocess.cc:304] Start cannot spawn child process: No such file or directory
2021-05-26 19:21:36.654664: F external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc:478] ptxas returned an error during compilation of ptx to sass: 'Internal: Failed to launch ptxas'  If the error message indicates that a file could not be written, please verify that sufficient filesystem space is provided.

[samuela]

I found that it's essential for the CUDA toolkit binaries to be in your $PATH. XLA_FLAGS seems to be unrelated to this particular issue.

[samuela]

Note to self: ptxas doesn't seem to need to be in the user's PATH as long as it's in the PATH for the right packages. Most recently I fixed this issue by:

• switching from tensorflow to tensorflowWithCuda
• sudo nix-channel --update
• sudo nixos-rebuild switch --upgrade
• restart
• build tensorflowWithCuda
  Not sure exactly what went wrong or what fixed it, but it's back to working now...

[carlthome]

Hmm, I'm running tensorflow==2.9.0rc2 inside a devcontainer from docker.io/nvidia/cuda and also got errors due to ptxas not being on PATH.

Changing base image from 11.4.0-cudnn8-runtime-ubuntu20.04 to 11.4.0-cudnn8-devel-ubuntu20.04 solved it for me (sensible in hindsight but mentioning it here for reference).

[jayagami]

install nvidia-cuda-toolkit or switch to devel both solved my problem.

[zehsilva]

Quickfix for the particular case of running into this problem when using conda to manage/install the packages in local environements.

• Install conda install -c nvidia cuda-nvcc (so that the ptxas binaries are added to your environment).

---

I had already installed cudatoolkit and a installation of jax that seemed to work, until I realized this error was there when trying to call jax.random.PRNGkey(42). After a bit of investigation of what was installed by the cudatoolkit, what is on $PATH variable and such, I realized the binaries was not present in the correct place (in this case the bin/ folder inside the environment I was working). Luckily installing the package cuda-nvcc solves it.

[maulberto3]

In my WSL2, I ran pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html (as per the docs) and then conda install -c nvidia cuda-nvcc (as this suggestion) and now, when importing jax, I have no more warnings about ptxas or different cuda versions.

[maulberto3]

Having said that, in my case, gpu is still slower than cpu:

[maulberto3]

Current Google Colab build seems to be slower too https://colab.research.google.com/drive/1szwLQW4OkEP-BjB5ekv1x16eDBD_Q4hy?hl=es#scrollTo=-4_BmfdcY-jX

[maulberto3]

So it seems that not all operations exhibit jax benefits, as doing other similar test, indeed gpu was faster, have a look at the same colab above for the updated version.

[Mahdi-Bashiri]

@zehsilva dear fellow , you have no idea that how much your guide/solution helped me out of my misery!

I love to thank you indeed.

P.S:
WSL on Windows10
Ubuntu distro. in WSL
TensorFlow v2.11 as in a tf conda env. (https://www.youtube.com/watch?v=0S81koZpwPA&t=22s) ===> ALSO SO HELPFUL !

While running a GAN model, using tensorflow and needing my RTX 3060 GPU, I had "kernel crashs" problem. It was about "Couldn't invoke ptxas --version" and so on.
I had checked almost all of the troubleshootings suggested by any source from the Internet. With no result and hope.

BUT, by your QuickFix tip, the whole problem was solved and now my code runs well.

Thanks a lot again.

[blue-codes-yep]

Maybe an add to this solution. By default, the command conda install -c nvidia cuda-nvcc installs the newest version of nvcc. As I am writing this, the version is 11.7 which need the latest driver 515 from nvidia. One can find the fitted version from here.

Thank you, had to do some other weird things like setting flags, and moving my nvvm folder but got everything working thanks!

[maulberto3]

In my WSL2, I ran pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html (as per the docs) and then conda install -c nvidia cuda-nvcc (as this suggestion) and now, when importing jax, I have no more warnings about ptxas or different cuda versions.

Having said this, for different reasons I had to reinstall locally everything, and I got an idea, before doing the conda install, i decided to hit nvidia-smi, turned out i got 12.2, clearly colliding with ptxas 12.3 for jax, so I figured would a quick easy GeForce driver update fix it towards 12.3? Turns out it did! So now, I hit nvidia-smi again and volila, I have 12.3, and ewhen going into jax, there is now no ptxas collindng issues!