How to correctly build from source with a cuda system install ?

[ccoulombe]

Hello,

Trying to correctly build jaxlib 0.4.34 from source with an already available cuda, cudnn, and nccl installation, it raises

external/xla/xla/tsl/cuda/cupti_stub.cc:16:10: fatal error: 'third_party/gpus/cuda/extras/CUPTI/include/cupti.h' file not found
   16 | #include "third_party/gpus/cuda/extras/CUPTI/include/cupti.h"

I could try to symlink this path (third_party/gpus/cuda/) to our actual installation ($EBROOTCUDA), or patch the header to use CPATH to correctly find the headers.

Commands used:

module load cuda/12.2 cudnn/9.2 nccl/2.18 python/3.11 clang/17
git clone --depth 1 https://github.com/jax-ml/jax.git -b release/0.4.34
cd jax/build
python build.py --target_cpu_features=release --enable_cuda --use_clang --cuda_compute_capabilities='8.0' --cuda_version='12.2' --cudnn_version='9.2.1' --bazel_options=--repo_env=LOCAL_CUDA_PATH=$EBROOTCUDA --bazel_options=--repo_env=LOCAL_CUDNN_PATH=$EBROOTCUDNN --bazel_options=--repo_env=LOCAL_NCCL_PATH=$EBROOTNCCL --bazel_options=--local_cpu_resources=$SLURM_CPUS_PER_TASK

How can I configure jaxlib to be built with the system installations ?

Thanks,

[hawkinsp]

Docs are here: https://github.com/openxla/xla/blob/main/docs/hermetic_cuda.md

However, note because of version compatibility, you shouldn't need to. Any reason you can't use the prebuilt jax artifacts? They should work with the versions you have.

[ccoulombe]

Thanks @hawkinsp, I missed this xla doc.

Usually we prefer to build from source, but I was running into runtime issues with the prebuilt jax/jaxlib (see #25344). I've now fixed the issue, and the prebuilt wheels works correctly.

Cheers