Thank you for your answer, I am sure your suggestion would have helped!
I finally understood the source of the problem. This network contains this layer: Linear(in_features=131072, out_features=1024, bias=True)
Which requires a huge number of gradients to be synchronized (~400 MB per iteration in fp32). As I don’t have any NVLink it makes sense that this synchronisation takes 500 ms I think.
So I ended up redesigning the network to avoid having such a huge fully connected layer.