UNKNOWN: Failed to determine best cudnn convolution algorithm

[jrabary]

Dear all,

Recently, my code fails with the following error:

UNKNOWN: Failed to determine best cudnn convolution algorithm: UNKNOWN: GetConvolveAlgorithms failed.

I work inside a container environment based on the tensorflow/tensorflow:latest-gpu docker image. Before the release of tensorflow 2.7, I used the tensorflow 2.6 image and every thing worked well. I switched to the 2.7 docker image since Jax GPU now requires python 3.8 at least.

[jrabary]

I finally make it works by installing cudnn 8.2. The container comes with cudnn 8.1.

[mattiasmar]

Do you think you changed something besides updating the cudnn version?
I'm working with the docker image nvidia/cuda:11.3.0-cudnn8-devel-ubuntu20.04 (CUDNN version 8.2.0.53) and I get the below error when running the example flax\examples\imagenet\configs\fake_data_benchmark.py.

RuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
%cudnn-conv = (f16[32,112,112,64]{2,1,3,0}, u8[0]{0}) custom-call(f16[32,224,224,3]{2,1,3,0} %copy.460, f16[7,7,3,64]{1,0,2,3} %convert.493), window={size=7x7 stride=2x2 pad=3_3x3_3}, dim_labels=b01f_01io->b01f, custom_call_target="__cudnn$convForward", metadata={op_type="conv_general_dilated" op_name="pmap(<unnamed wrapped function>)/conv_general_dilated[\n  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=(32, 224, 224, 3)\n  padding=((3, 3), (3, 3))\n  precision=None\n  preferred_element_type=None\n  rhs_dilation=(1, 1)\n  rhs_shape=(7, 7, 3, 64)\n  window_strides=(2, 2)\n]" source_file="/usr/local/lib/python3.8/dist-packages/flax/linen/linear.py" source_line=294}, backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"