Nvidia的NVLink加速了同一台机器上多个gpu之间的数据传输。我使用PyTorch在这样的机器上训练大型模型。
我明白为什么NVLink会使模型并行训练更快,因为通过一个模型将涉及几个gpu。
但是它会加速使用distributeddataparliel的数据并行训练过程吗?
如何在k
gpu上进行数据并行训练?
你把你的小批量分成k
部分,每个部分在不同的GPU上转发,并且在每个GPU上估计梯度。然而,(这是超级关键的)更新权重必须在所有gpu之间同步。这就是NVLink在数据并行训练中变得重要的地方。