That should be same as setting delay_reduce=True in apex right? That doesn’t fix it either.
I think the deadlock is between all_reduce calls in syncBNs’ backward() across the parallel heads and not between gradient all_reduce & syncBN all_reduce.
Removing the parallel head and having just 1 head works fine. Replacing syncBN with BN works fine for with parallel head model.
Thank you!