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

Question

在以下場景中，有 2 個 GPU，每個 GPU 的速度明顯不同：GPU0 比 GPU1 快約 40%。 我想將模型訓練 100k 步，通常如果速度相同，將在相當於 50k 步時達到。

但是，由於 GPU 的速度不同，當 GPU0 達到 50K 步時，GPU1 僅達到 30K 步。 實際上，該模型已經訓練了 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 完成后自動讓 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

Answer 1

由於 DDP 在每一步都完全同步梯度，因此較快的 GPU0 應始終等待較慢的 GPU1。

同步發生在后退步驟——allreduce 梯度。

DDP 不是為在異構環境中運行而設計的。 您可以考慮將輸入與兩個 GPU 的計算能力成正比，並讓 DDP 處理不均勻的 DDP 輸入。