Slow operation for convolution

[maxencefaldor]

Did anyone encounter the same issue updating to the last version of DreamerV3?

2024-05-01 12:03:30.285255: E external/xla/xla/service/slow_operation_alarm.cc:65] Trying algorithm eng28{k2=3,k3=0} for conv (bf16[16,5,5,32]{3,2,1,0}, u8[0]{0}) custom-call(bf16[16,67,67,1024]{3,2,1,0}, bf16[32,63,63,1024]{3,2,1,0}), window={size=63x63 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
2024-05-01 12:03:30.661234: E external/xla/xla/service/slow_operation_alarm.cc:133] The operation took 1.376059286s
Trying algorithm eng28{k2=3,k3=0} for conv (bf16[16,5,5,32]{3,2,1,0}, u8[0]{0}) custom-call(bf16[16,67,67,1024]{3,2,1,0}, bf16[32,63,63,1024]{3,2,1,0}), window={size=63x63 rhs_reversal=1x1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward", backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} is taking a while...
2024-05-01 12:03:32.943688: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 0: 57856, expected 17152
2024-05-01 12:03:32.943700: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 1: 58112, expected 17664
2024-05-01 12:03:32.943703: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 2: 58112, expected 16768
2024-05-01 12:03:32.943704: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 3: 57856, expected 16896
2024-05-01 12:03:32.943706: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 4: 58112, expected 17024
2024-05-01 12:03:32.943708: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 5: 58112, expected 16768
2024-05-01 12:03:32.943710: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 6: 57856, expected 17152
2024-05-01 12:03:32.943711: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 7: 57856, expected 17024
2024-05-01 12:03:32.943714: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 8: 58112, expected 16768
2024-05-01 12:03:32.943716: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 9: 57856, expected 17280
2024-05-01 12:03:32.943719: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:747] Results mismatch between different convolution algorithms. This is likely a bug/unexpected loss of precision in cudnn.
(bf16[48,5,5,32]{3,2,1,0}, u8[0]{0}) custom-call(bf16[1024,33,33,32]{3,2,1,0}, bf16[1024,16,16,48]{3,2,1,0}), window={size=5x5 stride=2x2 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convBackwardFilter", backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} for eng47{k0=9,k2=7,k5=4,k14=2} vs eng20{k2=7,k3=0}
2024-05-01 12:03:32.943723: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:307] Device: NVIDIA RTX 6000 Ada Generation
2024-05-01 12:03:32.943724: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:308] Platform: Compute Capability 8.9
2024-05-01 12:03:32.943727: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:309] Driver: 12040 (INVALID_ARGUMENT: expected %d.%d, %d.%d.%d, or %d.%d.%d.%d form for driver version; got "1")
2024-05-01 12:03:32.943729: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:310] Runtime: <undefined>
2024-05-01 12:03:32.943731: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:317] cudnn version: 8.9.2
2024-05-01 12:03:32.956910: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 0: 57856, expected 17152
2024-05-01 12:03:32.956914: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 1: 58112, expected 17664
2024-05-01 12:03:32.956916: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 2: 58112, expected 16768
2024-05-01 12:03:32.956918: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 3: 57856, expected 16896
2024-05-01 12:03:32.956919: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 4: 58112, expected 17024
2024-05-01 12:03:32.956921: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 5: 58112, expected 16768
2024-05-01 12:03:32.956923: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 6: 57856, expected 17152
2024-05-01 12:03:32.956924: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 7: 57856, expected 17024
2024-05-01 12:03:32.956926: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 8: 58112, expected 16768
2024-05-01 12:03:32.956928: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 9: 57856, expected 17280
2024-05-01 12:03:32.956931: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:747] Results mismatch between different convolution algorithms. This is likely a bug/unexpected loss of precision in cudnn.
(bf16[48,5,5,32]{3,2,1,0}, u8[0]{0}) custom-call(bf16[1024,33,33,32]{3,2,1,0}, bf16[1024,16,16,48]{3,2,1,0}), window={size=5x5 stride=2x2 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convBackwardFilter", backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} for eng47{k0=9,k2=7,k5=4,k14=2} vs eng20{k2=6,k3=0}
2024-05-01 12:03:32.956935: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:307] Device: NVIDIA RTX 6000 Ada Generation
2024-05-01 12:03:32.956936: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:308] Platform: Compute Capability 8.9
2024-05-01 12:03:32.956938: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:309] Driver: 12040 (INVALID_ARGUMENT: expected %d.%d, %d.%d.%d, or %d.%d.%d.%d form for driver version; got "1")
2024-05-01 12:03:32.956941: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:310] Runtime: <undefined>
2024-05-01 12:03:32.956943: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:317] cudnn version: 8.9.2
2024-05-01 12:03:33.030591: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 0: 2320, expected 1056
2024-05-01 12:03:33.030599: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 1: 2336, expected 1048
2024-05-01 12:03:33.030601: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 2: 2320, expected 1048
2024-05-01 12:03:33.030603: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 3: 2320, expected 1056
2024-05-01 12:03:33.030605: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 4: 2320, expected 1032
2024-05-01 12:03:33.030606: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 5: 2320, expected 1048
2024-05-01 12:03:33.030608: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 6: 2304, expected 1064
2024-05-01 12:03:33.030609: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 7: 2304, expected 1040
2024-05-01 12:03:33.030612: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 8: 2320, expected 1040
2024-05-01 12:03:33.030614: E external/xla/xla/service/gpu/buffer_comparator.cc:151] Difference at 9: 2304, expected 1048
2024-05-01 12:03:33.030619: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:747] Results mismatch between different convolution algorithms. This is likely a bug/unexpected loss of precision in cudnn.
(bf16[64,5,5,64]{3,2,1,0}, u8[0]{0}) custom-call(bf16[1024,9,9,64]{3,2,1,0}, bf16[1024,4,4,64]{3,2,1,0}), window={size=5x5 stride=2x2 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convBackwardFilter", backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"cudnn_conv_backend_config":{"conv_result_scale":1,"activation_mode":"kNone","side_input_scale":0,"leakyrelu_alpha":0},"force_earliest_schedule":false} for eng10{k2=5,k12=6,k13=1,k14=3,k15=0,k17=7,k18=1,k23=0} vs eng20{k2=6,k3=0}
2024-05-01 12:03:33.030622: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:307] Device: NVIDIA RTX 6000 Ada Generation
2024-05-01 12:03:33.030624: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:308] Platform: Compute Capability 8.9
2024-05-01 12:03:33.030626: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:309] Driver: 12040 (INVALID_ARGUMENT: expected %d.%d, %d.%d.%d, or %d.%d.%d.%d form for driver version; got "1")
2024-05-01 12:03:33.030627: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:310] Runtime: <undefined>
2024-05-01 12:03:33.030630: E external/xla/xla/service/gpu/conv_algorithm_picker.cc:317] cudnn version: 8.9.2

[swsychen]

Me too

[danijar]

I do get those as well sometimes. They are XLA/CUDA issues but I think they are usually just warnings. Does the code continue to run after compilation is done?

Otherwise, you can try adding --xla_gpu_autotune_level=0 to the XLA_FLAGS in dreamerv3/jaxagent.py or try updating JAX to the newest vesion.

[maxencefaldor]

Yes, the code continues to run but I wonder if it runs slower than it should. I am already using JAX last version (0.4.26).

I noticed a decreased in fps after this update but I am not sure if this is related to those errors/warnings.

edit: I didn't mean to close the issue, sorry about that.

[danijar]

The decreased FPS are probably because the model architecture is different now.

You can save some flops and memory by setting --'(enc|dec).simple.outer' False --'(enc|dec).simple.mults' 1,2,3,4 which gives very similar performance unless the images are visually more complex.

Closing this issue because those are XLA/CUDA warnings unrelated to Dreamer itself. I would also be happy if they got fixed.

[maxencefaldor]

For anyone interested, changing jax.compute_dtype from bfloat16 to float32 fixed the warnings for me.

However, using float32 is a little bit slower than bfloat16 going from 500 to 430 fps/policy on my hardware.