JAX CUDNN Unknown Status Error.

[pranavsubramani]

JAX error -- CUDNN Unknown status

I looked through some of the other issues but couldn't find anything particularly related to this. For reference, I installed CUDA, CUDNN, verified the install with the mnistCUDNN file. Tried nvcc -V which returned the following:

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Sun_Feb_14_21:12:58_PST_2021
Cuda compilation tools, release 11.2, V11.2.152
Build cuda_11.2.r11.2/compiler.29618528_0

I ran a convolutional neural network with per example gradients and the following was the error:

2021-04-02 17:39:23.593394: W external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.cc:780] Failed to determine best cudnn convolution algorithm: Internal: All algorithms tried for convolution %custom-call = (f32[50,14,14,16]{2,1,3,0}, u8[0]{0}) custom-call(f32[50,28,28,1]{2,1,3,0} %bitcast, f32[8,8,1,16]{1,0,2,3} %copy.4), window={size=8x8 stride=2x2 pad=3_3x3_3}, dim_labels=b01f_01io->b01f, custom_call_target="__cudnn$convForward", metadata={op_type="conv_general_dilated" op_name="conv_general_dilated[ batch_group_count=1\n                      dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 3, 1, 2), rhs_spec=(3, 2, 0, 1), out_spec=(0, 3, 1, 2))\n                      feature_group_count=1\n                      lhs_dilation=(1, 1)\n                      lhs_shape=(50, 28, 28, 1)\n                      padding=((3, 3), (3, 3))\n                      precision=None\n                      rhs_dilation=(1, 1)\n                      rhs_shape=(8, 8, 1, 16)\n                      window_strides=(2, 2) ]"}, backend_config="{\"algorithm\":\"0\",\"tensor_ops_enabled\":false,\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}" failed. Falling back to default algorithm. 

Convolution performance may be suboptimal.
2021-04-02 17:39:23.666103: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:1960] Execution of replica 0 failed: Unknown: CUDNN_STATUS_EXECUTION_FAILED
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(3296): 'cudnnConvolutionForward( cudnn.handle(), alpha, input_nd.handle(), input_data.opaque(), filter_nd.handle(), filter_data.opaque(), conv.handle(), ToConvForwardAlgo(algorithm_desc), scratch_memory.opaque(), scratch_memory.size(), beta, output_nd.handle(), output_data.opaque())'

The final traceback is as follows:

RuntimeError: Unknown: CUDNN_STATUS_EXECUTION_FAILED
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(3296): 'cudnnConvolutionForward( cudnn.handle(), alpha, input_nd.handle(), input_data.opaque(), filter_nd.handle(), filter_data.opaque(), conv.handle(), ToConvForwardAlgo(algorithm_desc), scratch_memory.opaque(), scratch_memory.size(), beta, output_nd.handle(), output_data.opaque())'

I checked my LD_LIBRARY_PATH and my PATH and they appear to have the right location present in them. This is also with CUDA 11.2 and the latest versions of JAX/jaxlib that are present in the library.

Would appreciate any insight into this :)

[AndrewCSQ]

I've had the best success running CUDA-accelerated JAX inside a tensorflow-gpu image. There, all of the library locations "just work" and are correct. See my discussion.

[josephrocca]

I'm seeing a similar CUDNN_STATUS_EXECUTION_FAILED error. Not sure if it's related to yours, but I've documented it here with a minimal docker reproduction: huggingface/transformers#13416

@AndrewCSQ I tried using your docker image, but it unfortunately the error still occurs.

EDIT: Solution here: huggingface/transformers#13416 (comment) Seems like the problem is that JAX tries to pre-allocate too much memory. I ran export XLA_PYTHON_CLIENT_MEM_FRACTION=.7 in the shell before running Python (as explained here) and that solved it.

[marikgoldstein]

this really helped me @josephrocca , thank you!!!

[TANSixu]

Ran into the same error, this solution works for me! Thank you!!!!!! @josephrocca

[Dikshuy]

thanks @josephrocca!!!