pin cuda-nvcc to get a better working jaxlib from conda-forge

[ngam]

Description

cc @hawkinsp

xref pangeo-data/pangeo-docker-images#387

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random

key = random.PRNGKey(0)

fails with

2022-10-12 15:16:32.114395: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_asm_compiler.cc:57] cuLinkAddData fails. This is usually caused by stale driver version.
2022-10-12 15:16:32.114450: E external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:1325] The CUDA linking API did not work. Please use XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1 to bypass it, but expect to get longer compilation time due to the lack of multi-threading.
...
XlaRuntimeError: UNKNOWN: no kernel image is available for execution on the device

with cuda-nvcc>11.6.*

What jax/jaxlib version are you using?

jaxlib 0.3.15 (cuda)

Which accelerator(s) are you using?

GPU

Additional system info

No response

NVIDIA GPU info

K80/T4 tested

[hawkinsp]

We have a planned fix for this, which boils down to: Look at the driver version. Is it less than the ptxas version? If so, disable parallel compilation.

The issue comes when your ptxas version is newer than the driver you have installed.

[ngam]

Thanks for the speedy response. Yep, that seems to be it (see xref pangeo-data/pangeo-docker-images#387).

[ngam]

import os
os.environ["XLA_FLAGS"]="--xla_gpu_force_compilation_parallelism=1"

also resolves this issue; closing!