@alekhka as a temporary workaround, you could try doing allreduce manually after the backward pass. (see this comment 17). This will be slower as there will be no overlapping between computation and communication, but hopefully can avoid the hang problem.