GPU环境下jax.random.uniform触发段错误，CPU及normal函数正常

问题描述 #

在bash终端执行以下4条测试命令：

JAX_PLATFORM_NAME=cpu python -c "import jax; import jax.numpy as jnp; key = jax.random.PRNGKey(1); print(jax.random.uniform(key, (2, 2),))"
JAX_PLATFORM_NAME=cpu python -c "import jax; import jax.numpy as jnp; key = jax.random.PRNGKey(1); print(jax.random.normal(key, (2, 2),))"
JAX_PLATFORM_NAME=gpu python -c "import jax; import jax.numpy as jnp; key = jax.random.PRNGKey(1); print(jax.random.uniform(key, (2, 2),))"
JAX_PLATFORM_NAME=gpu python -c "import jax; import jax.numpy as jnp; key = jax.random.PRNGKey(1); print(jax.random.normal(key, (2, 2),))"

其中仅GPU环境下的jax.random.uniform采样命令触发Segmentation fault (core dumped)，退出码为139（长程序中为245），其余命令均正常执行。

环境信息 #

NVIDIA驱动与CUDA版本 #

NVIDIA-SMI 565.57.01              Driver Version: 565.57.01      CUDA Version: 12.7    
NVIDIA GeForce RTX 4090  

初始JAX版本 #

Name: jax
Version: 0.4.37
Name: jaxlib
Version: 0.4.36

版本更新后的异常 #

将jax和jaxlib升级至0.6.2版本后，仍触发段错误，同时出现版本不兼容警告：

~/anaconda3/envs/unc/lib/python3.10/site-packages/jaxlib/plugin_support.py:71: RuntimeWarning: JAX plugin jax_cuda12_plugin version 0.4.33 is installed, but it is not compatible with the installed jaxlib version 0.6.2, so it will not be used.
  warnings.warn(
Segmentation fault (core dumped)

警告信息翻译：

~/anaconda3/envs/unc/lib/python3.10/site-packages/jaxlib/plugin_support.py:71: 运行时警告：已安装JAX插件jax_cuda12_plugin版本0.4.33，但该版本与当前jaxlib 0.6.2不兼容，因此不会被启用。

解决思路 #

• 清理不兼容的CUDA插件：卸载版本不匹配的jax_cuda12_plugin，执行命令：

  pip uninstall -y jax_cuda12_plugin
  

  随后重新安装与当前jaxlib版本匹配的CUDA支持包，推荐直接安装带CUDA 12.x支持的完整JAX包：

  pip install "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
  

• 严格匹配jax与jaxlib版本：初始环境中jax（0.4.37）与jaxlib（0.4.36）版本不一致，这本身会引发兼容性问题。升级或安装时必须保证两者版本完全相同。
• 验证CUDA环境完整性：检查CUDA Toolkit是否正确安装，确认nvcc --version输出的编译器版本与驱动适配；同时检查CUDA_PATH、LD_LIBRARY_PATH等环境变量是否配置正确。
• 降级至稳定兼容版本：若最新版本仍存在问题，可尝试降级到适配CUDA 12.7的稳定JAX版本（如0.4.28及以上支持CUDA 12.x的版本），避免新版本的适配bug。
• 检查GPU架构适配：RTX 4090属于Ampere架构，确保使用的JAX版本支持该架构（通常0.3.x及以上版本均支持，但旧版本可能存在例外）。

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