分布式入门，怎样用PyTorch实现多GPU分布式训练(3)

假设读者熟知标准的随机梯度下降算法(SGD)，该算法常用于训练深度学习模型。我们现在看到的是 SGD 的一个变体——同步 SGD(synchronous SGD)，它利用 all-reduce collective 来进行扩展。我们先从标准 SGD 的数学公式开始吧。

其中 D 是一个样本集合(mini-batch)，θ 是所有参数的集合，λ 是学习率，Loss(X, y) 是某个损失函数在 D 中所有样本上的均值。

同步 SGD 所依赖的核心技巧是将更新规则中的求和在更小的 (mini)batch 子集上进行分割。D 被分割成 R 个子集 D₁, D₂, . .(推荐每个子集具有相同数量的样本)，所以将标准的 SGD 更新公式中的求和进行分割，得到：

现在，因为梯度算子在求和算子上是分布式的，所以我们得到：

我们从中得到了什么?

看一下上面方程中单独的梯度项(方括号里面)。它们现在可以被独立地计算，然后加起来得到原始的梯度，而且没有任何损失/近似。这就是数据并行。下面是整个过程：

• 将整个数据集分成 R 个等大的数据块(子集)。这里的字母 R 代表的是 replica(副本)。
• 使用 MPI 启动 R 个进程/rank，将每个进程绑定到一个数据块上。
• 让每个 rank 使用大小为 B 的 mini-batch(dᵣ)(dᵣ来自该 rank 分配到的数据块 D_r)计算梯度，即 rank r 计算。
• 将所有 rank 的梯度进行求和，然后将得到的梯度对每个 rank 可见，再进行进一步处理。

最后一点就是 all-reduce 算法。所以，每次在所有 rank 使用大小为 B 的 mini-batch(dᵣ)计算完梯度以后，都必须执行 all-reduce。需要注意的一点是，将全部 R 个 rank(使用大小为 B 的 mini-batch 计算出)的梯度相加之后会得到一个有效的批大小：

下面是实现的关键部分(没有展示样板代码)：

model = LeNet() 
# first synchronization of initial weights 
sync_initial_weights(model, rank, world_size) 
 
optimoptimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.85) 
 
model.train() 
for epoch in range(1, epochs + 1): 
    for data, target in train_loader: 
        optimizer.zero_grad() 
        output = model(data) 
        loss = F.nll_loss(output, target) 
        loss.backward() 
 
        # The all-reduce on gradients 
        sync_gradients(model, rank, world_size) 
 
        optimizer.step() 
 
def sync_initial_weights(model, rank, world_size): 
    for param in model.parameters(): 
        if rank == 0: 
            # Rank 0 is sending it's own weight 
            # to all it's siblings (1 to world_size) 
            for sibling in range(1, world_size): 
                dist.send(param.data, dst=sibling) 
        else: 
            # Siblings must recieve the parameters 
            dist.recv(param.data, src=0) 
 
 
def sync_gradients(model, rank, world_size): 
    for param in model.parameters(): 
        dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM) 

• 全部 R 个 rank 都使用随机权重创建自己的模型副本。
• 单个具有随机权重的副本可能导致在初始的时候不同步。推荐在所有的副本上同步初始权重，sync_initial_weights(..) 就是在做这件事。让任何一个 rank 将自己的权重发送到它的兄弟 rank，兄弟 rank 必须接收这些权重并用来初始化它们自身。
• 从每个 rank 对应的数据部分取出一个 mini-batch(大小为 B)，计算前向和反向传递(梯度)。作为配置的一部分，这里需要重点注意的一点是：所有的进程/rank 应该让自己那部分数据可见(通常是在自己的硬盘上或者在共享文件系统中)。
• 把求和作为 reduction 运算，对每一个副本上的梯度执行 all-reduce 集体。sync_gradients(..) 会完成梯度同步。
• 梯度同步之后，每个副本能够在自己的权重上独立地执行标准的 SGD 更新。optimizer.step() 正常运行。

现在问题来了：我们如何确保独立的更新保持同步?

（编辑：ASP站长网）