"Each GPU has a different forward pass (in terms of used parameters)" slows down training

Hello,

I have a model where in each forward pass, layers are randomly skipped with some probability, something like this:

def forward(self, x):
    p = 0.2
    for i in range(self.n_layers):
        # drop this layer with probability p
        random_number = random.uniform(0, 1)
        if self.training and (random_number < p):
            continue
        x = self.layers[i](x)
    return x

Theoretically, higher p should result in faster training as the number of skipped layers is higher. However, when training this model using apex.parallel.DistributedDataParallel, I observed a very significant slowdown (when p > 0). I try timing the forward passes and indeed, they were faster when p is higher, so I think the issue lies in the backward or the gradient gathering steps.

Note that this issue does not occur when using torch.nn.parallel.DistributedDataParallel (I had to set find_unused_parameters = True for it to work): faster training for higher p.

Could you please help checking this?
When all layers are used (i.e. p = 0), apex was faster than torch.nn (especially on multiple nodes), so I hope that I can use apex for all my experiments.
Thank you very much in advance!

Start with apex.parallel.DistributedDataParallel and compare its backward and gradient-gathering behavior against torch.nn.parallel.DistributedDataParallel using the reported model, varying p between 0 and values greater than 0. Confirm whether skipped layers and unused parameters cause the slowdown, and consider the issue resolved when higher skip probabilities produce faster training without incorrect gradients.