概述
最近新学习的JAX在使用时，提示：
2024-05-26 19:46:32.016388: W external/xla/xla/service/gpu/nvptx_compiler.cc:760]
The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.5.40).
Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation.
You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.
该警告信息说明您的系统中NVIDIA驱动程序的CUDA版本（12.2）低于PTXAS（CUDA编译器）的版本（12.5.40）。PTXAS是将CUDA代码编译为PTX代码（一种中间表示形式）再进一步转换为GPU可执行代码的关键工具。
当PTXAS的版本高于NVIDIA驱动程序支持的CUDA版本时，可能会出现兼容性问题。这里，因为驱动程序较旧，XLA（Accelerated Linear Algebra，加速线性代数运算的组件）决定禁用并行编译以避免潜在的问题。并行编译能够显著加快编译速度，特别是在处理大型计算图时。因此，禁用并行编译可能会导致你的程序在编译阶段变慢。
升级 cuda 到 12.5版本
首先打开 nvdia 官网[https://developer.nvidia.com/cuda-downloads] 下载，我的是 ubuntu 24.04 但还没有对应的版本，我就选择了 22.04 一样安装成功了。
安装步骤：
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/12.5.0/local_installers/cuda-repo-ubuntu2204-12-5-local_12.5.0-555.42.02-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-12-5-local_12.5.0-555.42.02-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-12-5-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-5
sudo apt-get install -y cuda-drivers
注意最后安装 cuda-drivers 还可以选择安装 NVIDIA Driver kernel module, 直接嵌入到内核，但我没有安装成功，最后又回退到经典安装驱动。
小工具
- to install CUDA on Ubuntu, you must have the Nvidia drivers installed on Ubuntu. Run this command to get the recommended Nvidia driver for your device:
sudo ubuntu-drivers devices
- nvidia-smi