【Jax报错】CUDA backend failed to initialize: Unable to load cuPTI

如题，在本地跑xai的grok-1的时候遇见的问题。

首先你的cuda应该是安装好的，也就是bash下nvidia-smi可以显示、python下torch.cuda.is_available()返回True。

在执行

import jax

jax.local_device_count()

的时候报错这个。

在/usr/local/cuda/extras/CUPTI/lib64应该是包含so文件的，这时候在~/.bashrc加上export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/extras/CUPTI/lib64，然后source ~/.bashrc就行了。

import jax
jax.devices()

显示应该是[cuda(id=0), cuda(id=1), cuda(id=2), cuda(id=3), cuda(id=4), cuda(id=5), cuda(id=6), cuda(id=7)]这样的，有几个卡就显示几个。

如果没有显示，那要看你是不是安装的jax的gpu版本，参考JAX: 库安装和GPU使用，解决不能识别gpu问题