在NVIDIA DGX设备上运行JAX出现ptxas调用失败报错如何解决？

该错误本质是JAX依赖的XLA GPU编译器无法调用NVIDIA CUDA Toolkit自带的PTX汇编工具ptxas，无法将PTX代码编译为GPU可执行的SASS机器码，最终触发核心转储退出。

原始报错信息：

2021-10-25 13:00:05.863667: W external/org_tensorflow/tensorflow/stream_executor/gpu/asm_compiler.cc:80] Couldn't get ptxas version string: INTERNAL: Couldn't invoke ptxas --version
2021-10-25 13:00:05.864713: F external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc:435] 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.
Aborted (core dumped)

• 执行ptxas --version命令，如果返回命令不存在，说明CUDA Toolkit未安装，或者环境变量未正确配置
• 执行echo $PATH检查是否包含CUDA的bin目录，默认系统级CUDA的bin路径为/usr/local/cuda/bin
• 执行ls -l $(which ptxas)确认当前用户对ptxas文件有可执行权限
• 执行df -h /tmp检查临时目录剩余空间，至少保留1GB以上空闲空间用于编译临时文件存储
• 执行pip show jax jaxlib查看当前安装的JAX版本，确认和服务器CUDA版本匹配

• 无需修改系统CUDA配置的快速方案：直接安装对应CUDA版本的预编译jaxlib，CUDA 11.x版本执行pip install "jax[cuda11_pip]"，CUDA 12.x版本执行pip install "jax[cuda12_pip]"，安装时指定JAX官方的CUDA镜像源即可
• 系统已安装CUDA的修复方案：在运行程序前执行以下环境变量配置命令，也可以将命令写入~/.bashrc永久生效：

  export PATH=/usr/local/cuda/bin:$PATH
  export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
  

• 无系统CUDA操作权限的方案：使用conda在个人用户目录安装独立的CUDA Toolkit，从NVIDIA官方conda源安装即可，安装完成后conda会自动配置对应环境变量
• 磁盘空间不足的方案：执行export TMPDIR=/path/with/enough/space将临时目录指定到空闲空间充足的分区后再运行程序

内容的提问来源于stack exchange，提问作者Igor Rivin