WSL Ubuntu虚拟机中启用CUDA的JAX安装运行报错问题

在WSL虚拟机中尝试用GPU运行JAX，按官方文档执行命令：

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

先后在Miniconda环境、全新Ubuntu WSL的干净venv中尝试，手动安装JAX、CUDA和Cudnn也出现相同错误。主机PC的NVIDIA驱动已更新，GPU为GTX 1070Ti。

运行测试代码：

from jax import random
key = random.PRNGKey(0)

出现如下错误：

2023-04-16 22:41:14.651263: W external/xla/xla/stream_executor/cuda/cuda_dnn.cc:397] There was an error before creating cudnn handle: cudaErrorNotSupported : operation not supported
2023-04-16 22:41:15.925044: E external/xla/xla/stream_executor/cuda/cuda_driver.cc:1207] Failed to get stream capture info: operation not supported
2023-04-16 22:41:16.034579: E external/xla/xla/pjrt/pjrt_stream_executor_client.cc:2432] Execution of replica 0 failed: INVALID_ARGUMENT: stream is uninitialized or in an error state
---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
/root/ML/JaxTest.py in line 2
      1 from jax import random
----> 2 key = random.PRNGKey(0)

File ~/ML/MLVenv/lib/python3.10/site-packages/jax/_src/random.py:136, in PRNGKey(seed)
    133 if np.ndim(seed):
    134   raise TypeError("PRNGKey accepts a scalar seed, but was given an array of"
    135                   f"shape {np.shape(seed)} != (). Use jax.vmap for batching")
--> 136 key = prng.seed_with_impl(impl, seed)
    137 return _return_prng_keys(True, key)

File ~/ML/MLVenv/lib/python3.10/site-packages/jax/_src/prng.py:270, in seed_with_impl(impl, seed)
    269 def seed_with_impl(impl: PRNGImpl, seed: Union[int, Array]) -> PRNGKeyArray:
--> 270   return random_seed(seed, impl=impl)

File ~/ML/MLVenv/lib/python3.10/site-packages/jax/_src/prng.py:561, in random_seed(seeds, impl)
    559 else:
    560   seeds_arr = jnp.asarray(seeds)
--> 561 return random_seed_p.bind(seeds_arr, impl=impl)

File ~/ML/MLVenv/lib/python3.10/site-packages/jax/_src/core.py:360, in Primitive.bind(self, *args, **params)
    357 def bind(self, *args, **params):
    358   assert (not config.jax_enable_checks or
    359           all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args
--> 360   return self.bind_with_trace(find_top_trace(args), args, params)

File ~/ML/MLVenv/lib/python3.10/site-packages/jax/_src/core.py:363, in Primitive.bind_with_trace(self, trace, args, params)
    362 def bind_with_trace(self, trace, args, params):
--> 363   out = trace.process_primitive(self, map(trace.full_raise, args), params)
    364   return map(full_lower, out) if self.multiple_results else full_lower(out)

File ~/ML/MLVenv/lib/python3.10/site-packages/jax/_src/core.py:817, in EvalTrace.process_primitive(self, primitive, tracers, params)
    816 def process_primitive(self, primitive, tracers, params):
--> 817   return primitive.impl(*tracers, **params)

File ~/ML/MLVenv/lib/python3.10/site-packages/jax/_src/prng.py:573, in random_seed_impl(seeds, impl)
    571 @random_seed_p.def_impl
    572 def random_seed_impl(seeds, *, impl):
--> 573   base_arr = random_seed_impl_base(seeds, impl=impl)
    574   return PRNGKeyArray(impl, base_arr)

File ~/ML/MLVenv/lib/python3.10/site-packages/jax/_src/prng.py:578, in random_seed_impl_base(seeds, impl)
    576 def random_seed_impl_base(seeds, *, impl):
    577   seed = iterated_vmap_unary(seeds.ndim, impl.seed)
--> 578   return seed(seeds)

File ~/ML/MLVenv/lib/python3.10/site-packages/jax/_src/prng.py:817, in threefry_seed(seed)
    812 k1 = convert(
    813     lax.shift_right_logical(seed, lax_internal._const(seed, 32)))
    814 with jax.numpy_dtype_promotion('standard'):
    815   # TODO(jakevdp): in X64 mode, this can generate 64-bit computations for 32-bit
    816   # inputs. We should avoid this.
--> 817   k2 = convert(jnp.bitwise_and(seed, np.uint32(0xFFFFFFFF)))
    818 return lax.concatenate([k1, k2], 0)

    [... skipping hidden 8 frame]

File ~/ML/MLVenv/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py:1916, in ExecuteReplicated.__call__(self, *args)
   1911   self._handle_token_bufs(
   1912       results.disassemble_prefix_into_single_device_arrays(
   1913           len(self.ordered_effects)),
   1914       results.consume_token())
   1915 else:
-> 1916   results = self.xla_executable.execute_sharded(input_bufs)
   1917 if dispatch.needs_check_special():
   1918   out_arrays = results.disassemble_into_single_device_arrays()

XlaRuntimeError: INVALID_ARGUMENT: stream is uninitialized or in an error state

尝试安装CUDA 11版本时，错误类似，仅缺少W external/xla/xla/stream_executor/cuda/cuda_dnn.cc:397]这条警告。

• GPU架构兼容性限制：GTX 1070Ti属于Pascal架构（sm_61），新版JAX搭配CUDA 12时，默认不再对该旧架构提供完整支持。CUDA 12移除了部分针对Pascal的特性，导致初始化cuDNN句柄时触发cudaErrorNotSupported错误。
• WSL2环境的CUDA特性支持不足：WSL2对旧GPU的CUDA高级特性支持存在局限，比如错误中提到的流捕获（stream capture）功能，Pascal架构本身不支持该特性，进而导致获取流捕获信息失败，使CUDA流处于错误状态。
• JAX默认设备执行逻辑冲突：JAX默认会优先尝试用GPU执行运算，但旧架构无法满足所需的CUDA操作要求，导致流初始化失败，最终引发XlaRuntimeError。

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