None of the algorithms provided by cuDNN heuristics worked; trying fallback algorithms.

[JGameCreation]

Hi, I'm trying to do some more involved operations on my CNN gradients but I always get error messages similar to these:

2021-12-29 18:54:35.948861: W external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.cc:717] None of the algorithms provided by cuDNN heuristics worked; trying fallback algorithms.  Conv: (f32[8,32,32,3]{2,1,3,0}, u8[0]{0}) custom-call(f32[8,32,32,32]{2,1,3,0}, f32[3,3,3,32]{1,0,2,3}), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target="__cudnn$convBackwardInput", backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"
2021-12-29 18:54:36.037656: W external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.cc:717] None of the algorithms provided by cuDNN heuristics worked; trying fallback algorithms.  Conv: (f32[4,4,32,256]{1,0,3,2}, u8[0]{0}) custom-call(f32[8,19,19,32]{2,1,0,3}, f32[1,16,16,256]{2,1,0,3}), window={size=16x16}, dim_labels=f01b_i01o->01bf, feature_group_count=8, custom_call_target="__cudnn$convForward", backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"
2021-12-29 18:54:36.278243: W external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.cc:717] None of the algorithms provided by cuDNN heuristics worked; trying fallback algorithms.  Conv: (f32[8,8,8,32]{2,1,0,3}, u8[0]{0}) custom-call(f32[15,15,32,1024]{1,0,3,2}, f32[8,8,8,128]{2,1,3,0}), window={size=8x8}, dim_labels=01bf_o01i->f01b, feature_group_count=8, custom_call_target="__cudnn$convForward", backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"

Although these are only warnings and the calculation can still be carried out, the code seems to run quite slowly and in a different case the program even silently crashes after these messages which might be related. Is there anything I can do to make my convulations more amenable to GPU acceleration?
Best, Jannis

[matthias-wright]

Hi @JGameCreation, this looks like a compatibility problem to me. Which versions (JAX, CUDA, cuDNN) are you using? Maybe try running your code on Colab and see if the warnings still occur.

[matthias-wright]

Could you post your Colab here?

[JGameCreation]

Sure, here it is: https://colab.research.google.com/drive/1YJcplFuv2sAuAM3DTg2WHQmmP5X8WEzI?usp=sharing
What I really care about is the time it takes to calculate the individual Hessians wrt. perturbation parameters of the magnitude of the gradient wrt. network parameters. If you run it you will see that for a batch size of 1 it takes 1.4 seconds and for a batch size of 8 it takes 17.4 seconds.

[matthias-wright]

I think the part where you compute the Hessians in a loop 8/32 times rather than using vmap might be missing?

[JGameCreation]

Well, I assume that computing it in a loop 8 times is 8 times slower than a batch size of 1, so it should be about 11.2 seconds?

[PythonNut]

I wonder if the problem has something to do w/ the combination of jax 0.2.26 and jaxlib 0.1.75+cuda11.cudnn82

If you upgrade the packages in the colab to those versions, you get:

RuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm ...

Which seems related to the original error.