"Disabling cuDNN frontend for the following convolution" error

[samuela]

I'm seeing a bunch of warnings like

2022-02-05 02:42:54.851077: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc:4584] Disabling cuDNN frontend for the following convolution:
  input: {count: 6 feature_map_count: 3 spatial: 32 32  value_min: 0.000000 value_max: 0.000000 layout: BatchDepthYX}
  filter: {output_feature_map_count: 8 input_feature_map_count: 3 layout: OutputInputYX shape: 3 3 }
  {zero_padding: 1 1  pad_alignment: default filter_strides: 1 1  dilation_rates: 1 1 }
  ... because the current cuDNN version does not support it.

I'm able to reproduce this with a tiny batch size and number of channels, so I'm fairly certain I'm not running out of GPU memory. I'm running on a g4dn.xlarge EC2 instance with CUDA 11.5 and driver version 495.44.

Does anyone have any idea what could be causing this? How do I fix it? Happy to provide a thorough reproduction if necessary.

[samuela]

I was running cuDNN 8.1.1 and using the jaxlib cudnn805 wheel variant. Upgrading cuDNN to 8.3.0 and switching to the jaxlib cudnn82 wheel variant resolved the issue for me.

Relevant PRs/commits:

• cudnn: init cudnn_8_3 at 8.3.0 NixOS/nixpkgs#158218
• python3Packages.jaxlib-bin: 0.1.71 -> 0.1.75 NixOS/nixpkgs#158186
• samuela/nixpkgs@214f1a4