问题描述
Nvidia 的 NVLink 可加速同一台机器上多个 GPU 之间的数据传输。 我使用 PyTorch 在这样的机器上训练大型模型。
我明白为什么 NVLink 会使 model-parallel 训练更快,因为一次通过模型将涉及多个 GPU。
但是它会使用 DistributedDataParallel 加速数据并行训练过程吗?
解决方法
如何在 k
GPU 上进行数据并行训练?
您将小批量拆分为 k
部分,每个部分在不同的 GPU 上转发,并在每个 GPU 上估计梯度。但是,(这是非常重要的)更新权重必须在所有 GPU 之间同步。这也是 NVLink 对数据并行训练变得重要的地方。