|
假设读者熟知标准的随机梯度下降算法(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() 正常运行。
现在问题来了：我们如何确保独立的更新保持同步? （编辑：核心网）
【声明】本站内容均来自网络，其相关言论仅代表作者个人观点，不代表本站立场。若无意侵犯到您的权利，请及时与联系站长删除相关内容!
|
