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

问题描述

在以下情况下，有2个GPU，每个GPU的速度明显不同：GPU0比GPU1快40％。我想对模型进行10万步训练，如果速度相同，通常相当于5万步。

但是，由于GPU的速度不同，因此当GPU0达到5万步时，GPU1仅达到3万步。有效地，该模型已经训练了80k步。

在实践中，PyTorch的DistributedDataParallel是否可以与不同速度的GPU配合使用？当前，脚本[A]运行时，GPU0以其原始速度无需等待就可以继续运行GPU1，因此我想知道如何进行任何同步。我在同一步骤中在每个GPU中打印了模型的参数，它们的确确实存在很大差异。如果是这样，同步在哪里工作？

在DDP的原始源代码[B]中，似乎确实在模型的每个前向传递之前进行了同步。但是我不知道为什么会出现这种情况，每个GPU中的参数总和减少了总值的1-2％。

获取参数值的函数很简单：

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

    return total

有没有一种方法可以让GPU0在完成后自动接管GPU1的一些“剩余”训练？

[A]可以在此处找到正在运行的脚本： 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

解决方法

由于 DDP 在每一步都完全同步梯度，因此速度较快的 GPU0 应始终等待较慢的 GPU1。

同步发生在后退步骤——allreduce 梯度。

DDP 不是为在异构环境中运行而设计的。您可以考虑将输入与两个 GPU 的计算能力成正比，然后 let DDP handle uneven DDP inputs。