具有不同 GPU 速度的 PyTorch DistributedDataParallel 是否同步权重？

Question

In the following scenario, there are 2 GPUs, each with a significantly different speed: GPU0 is faster than GPU1 by around 40%. I want to train the the model for 100k steps, which ordinarily if the speeds are the same, would be reached in the equivalent of 50k steps.

However, since the GPUs are of different speeds, when GPU0 hits 50K steps, GPU1 reaches only 30K steps. Effectively, the model has been trained for 80k steps.

In practice, will PyTorch's DistributedDataParallel work with GPUs of different speed? Currently, the script [A] runs such that GPU0 proceeds at its original speed without waiting for GPU1, so I was wondering how any syncing would work. I printed the parameters of the model in each GPU at the same step, and they are indeed different significantly. If so, where does the syncing work?

In the original source code [B] for DDP, it does seem that syncing is made before each forward pass of the model. But I don't know why if this is the case, the sum of the parameters in each GPU is off by around 1-2% of the total value.

The function to get the parameters value is simply this:

def get_params_sum(net):
    total = 0
    for param in net.parameters():
        total = total + torch.sum(param)

    return total

Is there a way to automatically let GPU0 take over some "leftover" training of GPU1 when it's done?

[A] A running script can be found here: https://github.com/yangkky/distributed_tutorial/blob/master/src/mnist-distributed.py

[B] https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/distributed.py#L707

Answer 1

Since DDP fully syncs gradients at every step, the faster GPU0 should always wait for the slower GPU1.

The sync occurs at the backward step -- allreduce gradients.

DDP is not designed to run in a heterogeneous environment. You may consider dividing the input proportional to the compute power of the two GPUs, and let DDP handle uneven DDP inputs .