Potential tricky bug in the implementation of stochastic depth

🐛 Describe the bug

First of all sorry if this is not a bug. Though I think it is.

In training with stochastic depth, each forward pass randomly drop layers. In the current implementation, this doesn't seem to be the case when training on multiple GPUs, because the survival masks can be different across the GPUs.

For example, suppose that we have 4 residual layers trained on two GPUs, and that in the current iteration the survival masks are as follows:

GPU 0: True False True False
GPU 1: False True False True

In this case, all the layers have actually been used in the forward pass (that is, none of the layers have been dropped). It seems what happens here is equivalent to scaling the gradient by 0.5 (due to the reduction across the GPUs). As a result, the obtained algorithm is different from stochastic depth (even if it works very well).

If I am correct then a lot of existing implementations have the same issue. Adding mask synchronisation across the GPUs could be an easy fix.

Please let me know what you think. Thank you very much in advance!

Versions

Any version.

Start with torchvision/ops/stochastic_depth.py at the linked implementation and inspect how its survival masks are generated during multi-GPU training. Reproduce the reported behavior across two GPUs and determine the expected mask semantics; done means the implementation preserves stochastic-depth behavior across devices with regression coverage for the case described.