大规模 MLOps 工程（三）（3）

大规模 MLOps 工程（三）（2）https://developer.aliyun.com/article/1517788

8.2 参数服务器方法的梯度累积

本节介绍了基于参数服务器的分布式梯度下降的实现，并解释了梯度累积在实现中的作用。本节澄清了参数服务器方法的局限性，并激发了更高效的基于环的实现。

像 TensorFlow 1.x 这样的传统机器学习框架普及了基于参数服务器的方法，以在集群的多个节点之间分布梯度下降。 图 8.2 中描绘的参数服务器方法易于理解和实现。

图 8.2 梯度下降在工作节点和参数服务器之间进行分布以支持扩展

在图中，每个工作节点（使用虚线表示）根据训练数据集的单个分片执行梯度下降的前向和后向步骤（例如清单 8.4 中的内部循环中的步骤），以计算损失函数的分片特定梯度。请注意，在图 8.2 中，梯度具有与用于计算梯度的分片的下标对应的下标，就像图 8.1 中一样。

一旦工作节点计算出其梯度，它就将梯度发送到参数服务器（或参数服务器集群）进行处理。参数服务器（图 8.2 的右侧）等待累积来自工作节点的梯度，并使用累积的梯度执行梯度下降的优化步骤，计算下一次梯度下降的模型参数。然后，基于新计算的模型参数（在图 8.2 中表示为 w’），将下一个版本的模型发送到工作节点，取代以前的模型参数（在图 8.2 中表示为 w），确保每个节点使用相同和更新的模型的下一个梯度下降迭代。

图 8.2 中的分布式梯度下降的参数服务器实现是一种分布式数据并行（在本章的介绍中定义）方法。在分布式数据并行方法中，训练数据集被划分为独立且互不重复的子集，以便训练数据集分片和工作节点之间存在一对一的关系。接下来，每个工作节点使用一个分片和一个相同的模型参数副本计算梯度。

与替代的分布式数据并行方法（在本章的其余部分中讲解）不同，分布式梯度下降的参数服务器实现存在重大的可伸缩性问题：工作节点和参数服务器之间的网络连通性是通信瓶颈。具体而言，在实现的两个通信阶段中都存在通信带宽受限的问题：在从工作节点到参数服务器的梯度的多到一（或多到少）通信期间，以及在从参数服务器（多个参数服务器）到工作节点的更新模型参数的一到多（或少到多）通信期间。

8.3 引入逻辑环形梯度下降

本节介绍了在逻辑环网络中通信的节点的基本概念。本节不是为了提供实际节点并使其通过网络通信，而是使用在单节点环境中运行的简单 Python 程序来解释网络概念。一旦你对概念有了牢固的掌握，你将把它们应用到更复杂的、分布式的、多节点环境中。

与依赖于集中式参数服务器集群（第 8.2 节中所示方法）相反，基于逻辑环的分布式数据并行算法（例如 Horovod；github.com/horovod/horovod）避免了一对多和多对一通信的瓶颈，并且仅依赖于在环中与两个逻辑邻居通信的节点：前趋节点和后继节点。

图 8.3（左侧）的图示显示了四个节点，每个节点使用虚线表示，并表示为节点n[0]到n[3]，这些节点组织在一个逻辑环中。请注意，在公共云环境中的当代虚拟网络中，节点不必物理连接到彼此形成环：标准以太网网络足够。但是，在图中显示的逻辑环网络中，每个节点都受到限制，仅与其前趋节点和后继节点通信。正如您将在第 8.4 节中了解到的那样，这有助于限制分布式梯度下降每次迭代所需的网络带宽。

图 8.3 逻辑网络环（左）使用示例值解释

对于具有标识符n[i]的节点，后继节点的标识符被定义为n[(i+1) %] NODES，其中 NODES 是逻辑环中节点的总数。模运算确保通信模式形成一个环，其中具有最高标识符（始终为n[NODES-1]）的节点与标识符为 0 的节点进行通信，反之亦然。在环形网络中，如本章所述，每个节点只向后继节点发送数据。

使用类似的逻辑，基于环形网络的前趋节点的标识符被定义为n[(i-1) %] NODES，以便节点 0 可以与节点 1 和具有最高标识符值（NODES - 1）的节点进行通信。本章使用的环网络中的每个节点只从前趋节点接收数据。

就像第 8.2 节中解释的基于参数服务器的方法一样，图 8.3 中的节点处理训练数据集的独立碎片，以便g0 表示由节点n[0]计算的具有索引 0 的碎片的梯度值。继续使用第 8.2 节的示例，如果[0:250]0 是四个碎片中的第一个碎片，那么g0 表示由节点n[0]计算的第一个碎片的梯度值，使用模型参数值 w。因此，就像基于参数服务器的方法一样，基于环的方法也是数据并行分布式的。

在基于环的分布式数据并行实现中，不存在专用的参数服务器。相反，在集群中的每个节点完成梯度下降迭代的前向和后向步骤后，节点在逻辑环网络中通信，以便所有碎片的梯度在环中的每个节点上累积。

需要在节点之间通信什么样的信息，以确保模型参数值和累积梯度值完全同步和相同？在逻辑环中，由于每个节点只能向后继节点发送数据，因此节点只能从前驱节点的一系列梯度发送/接收操作中接收累积梯度。例如，为了让节点n[0]从节点n[1]到n[3]（图 8.4 的最右侧）累积梯度，需要三次迭代的发送/接收操作。这三次迭代从图 8.4 的左侧到右侧按顺序显示。正如您将在本章中观察到的那样，在由 NODES 个节点组成的多节点集群中，需要（NODES - 1）次发送/接收通信迭代。

图 8.4 在一个由四个节点组成的环中将梯度（求和）减少到节点 0，这是一个分布梯度的全局规约算法，用于在节点之间分发梯度。

列表 8.5 中的源代码提供了图 8.4 描述的逻辑的 Python 伪代码实现。在实现中，使用了 NODES 变量，该变量是使用训练数据集中训练示例的数量（常量 TRAINING_DATASET_SIZE 的值）与多节点集群中一个节点的内存中适合的训练示例的数量（IN_MEMORY_SHARD_SIZE 的值）之间的关系定义的。使用地板除法运算符//以确保 NODES 常量的值被设置为整数值，因为它稍后将用作 Python 范围操作的参数。

列表 8.5 Python 伪代码，以说明梯度减少到节点 0

NODES = \
  TRAINING_DATASET_SIZE // IN_MEMORY_SHARD_SIZE       ❶
GRADIENTS = [5., 3., 2., 1.]                          ❷
node_to_gradients = \
  dict(zip(range(NODES), GRADIENTS))                  ❸
for iter in range(NODES - 1):                         ❹
  node = (iter + 1) % NODES                           ❺
  grad = node_to_gradients[node]                      ❻
  next_node = (node + 1) % NODES                      ❼
  # simulate "sending" of the gradient value
  # over the network to the next node in the ring
  node_to_gradients[next_node] += grad                ❽

❶ 计算训练数据集所需的节点数量。

❷ 为演示分配任意的 GRADIENT 值，每个节点一个。

❸ 创建一个字典来跟踪节点计算的梯度。

❹ 执行 NODES - 1 次通信迭代。

❺ 从节点 iter+1 开始，以便在 NODES-1 后…

❻……迭代，节点 0 累积梯度。

❼下一个节点的标识符结束了环。

❽在节点对梯度进行累积。

一旦代码执行完毕，打印 node_to_gradients 字典的值。

print(node_to_gradients)

输出结果：

{0: 11.0, 1: 3.0, 2: 5.0, 3: 6.0}

其中键 0 对应于预期梯度，计算的 n[0]，值为 11，基于累积梯度 5+3+2+1。此外，请注意，由于图 8.4 不包括对n[0]以外的任何节点的梯度累积，因此 n[1]到n[3]的梯度保持不变。即将介绍的部分将解释如何确保在环中的所有节点上累积相同的梯度。

在三（节点-1）次迭代的第一次（在图 8.4 中以基于零的索引显示为迭代 0）中，节点n[1]发送并且节点n[2]接收节点n[1]在开始迭代 0 之前计算的梯度值 g1。由于在环中的通信目的是为了到达累积梯度，一旦接收 g1 梯度值，n[2]节点可以直接累积（加到）梯度值，以确保重用内存来存储累积梯度值：g1+g2。例如，如果每个节点上的每个梯度张量都是 400 MB，那么在环中的节点之间传输 400 MB 的数据，并且每个节点消耗 400 MB 的内存来存储累积梯度值。到迭代 0 结束时，节点n[2]累积了添加（即使用求和操作减少的）梯度。

因此，在第二次迭代（在图 8.4 中标记为迭代 1）期间，累积的梯度值从节点n[2]发送到节点n[3]，导致在第二次迭代结束时在节点 n[3]上累积的梯度值g1+g2+g3。

在这个示例中的最后一次迭代（在图 8.4 中标记为迭代 2）完成了对节点n[0]上的梯度的累积，将在这次迭代中计算的节点n[0]上的梯度* g0 加到从n* [3]收到的累积梯度上。由此得到的梯度，包括g0+g1+g2+g3，足以让n[0]计算出下一个优化步骤的模型参数值，这个步骤是由集群中的每个节点执行的梯度下降过程。

虽然图 8.4 和列表 8.5 中示例的三次迭代实现了梯度的累积（reduce 步骤）到单个节点，但要使分布式数据并行梯度下降工作，环中的每个节点必须访问整个累积梯度：g0 + g1 + g2 + g3。除非每个节点都可以访问累积梯度，否则节点无法执行使用累积梯度更改模型参数值的梯度下降步骤。即将介绍的各节将基于列表 8.5 中的 reduce 步骤来解释整个分布式梯度下降算法的 reduce-all 阶段。

8.4 理解基于环形的分布式梯度下降

虽然第 8.3 节描述的天真的基于环的 reduce 操作可以消除对参数服务器的需求，并确保梯度值在环形多节点集群中的各个计算节点上被减少（累积），但它存在一些缺点。随着训练数据集的增长（这是可以预料的），集群中的节点数量必须增长以跟上。这也意味着集群需要的总带宽必须随着节点数量的增加而增长，因为每个节点在每次迭代期间都必须将整个梯度发送给环中的下一个节点。在本节中，您将了解基于环形分布式数据并行算法（例如，著名的 Horovod 算法）如何在规模化情况下帮助有效利用带宽，其中训练节点的数量和训练示例的数量都增加。

Horovod 算法可以支持训练数据集的增长（以及集群中节点的数量），同时保持带宽需求恒定或甚至降低带宽要求。为了支持这一点，Horovod 依赖于两个分离且独立的环形通信阶段：（1）reduce-scatter 和（2）all-gather。在两个阶段中，Horovod 不是在节点之间发送/接收整个梯度数据，而是只通信梯度的一个单一段落，其中默认情况下段落的大小是梯度大小乘以 ，其中NODES是环集群中的工作节点数。因此，增加工作节点的数量以与训练数据集大小成比例地减少节点间通信的带宽需求。

那么梯度的 段 是什么？你可以把每个段视为梯度的逻辑分区，如图 8.5 所示。在图中，节点 n[0] 计算的梯度 g[0]，基于训练数据集分片 [0:250][0]（其中 [0:250] 是 Python 切片表示法），依次被分成 NODES 段，以便默认情况下，每个段都存在大致相等数量的梯度值。继续之前的例子，梯度占据了 400 MB 的数据量（例如模型参数的 32 位浮点梯度值的每 100,000,000 个字节为 4 字节），每个段是 100 MB 的互斥逻辑分区的相同模型张量。请注意，在这种情况下，由于分片是由节点 n[0] 计算的，因此每个四个段中的 i 都使用 s[i](n[0]) 进行注释。

图 8.5 Horovod 用于节点间通信的梯度段

还要注意，虽然在图 8.5 的框架的水平轴上不能累积（相加）段，但是可以沿垂直轴累积段。此外，图 8.5 段框架下方显示的段 s[i] 对应于每个节点计算的相应段的累积。例如，s[0] 等于 s[0](n[0]) + s[1](n[1]) + s[2](n[2]) + s[3](n[3])。因此，图 8.5 框架下方显示的 s[0]s[1]s[2]s[3] 段等同于将累积梯度 g[0] + g[1] + g[2] + g[3] 逻辑分割为段所需以执行梯度下降的优化步骤。

就像在列表 8.5 中介绍的基于环形减少步骤一样，本章的其余部分使用 Python 伪代码来解释 Horovod 算法。请回忆，对于一个分布式数据并行算法（如 Horovod）要正确工作，环形集群中的每个节点必须初始化具有模型参数的相同副本。在列表 8.6 中，Python 张量列表 W 被用来表示相同的模型。请注意，W 中的每个张量都是使用来自 w_src 的值初始化的，w_src 是从标准正态分布中抽样的伪随机值张量。

列表 8.6 W 存储模型张量的相同副本

pt.manual_seed(42)
w_src = pt.randn((4,))
W = [pt.tensor(w_src.detach().numpy(),
                requires_grad=True) for _ in range(NODES)]

为了重复使用列表 8.4 中的训练数据集张量 X_train 和 y_train，Horovod 算法的以下解释创建了一个 PyTorch DataLoader，它将训练数据集分成每个 IN_MEMORY_SHARD_SIZE 记录的分片。不要被列表 8.7 中 DataLoader 的 batch_size 参数所迷惑；虽然该参数用于分割源 TensorDataset，但是单个分片不会作为批量用于更新模型的参数。

列表 8.7 使用 PyTorch DataLoader 进行分片的梯度下降步骤

from torch.utils.data import TensorDataset, DataLoader
train_dl = DataLoader(TensorDataset(y_train, X_train), \
                      batch_size = IN_MEMORY_SHARD_SIZE,
                      shuffle = False)
for node, (y_shard, X_shard) in zip(range(NODES), train_dl):
  y_est = forward(W[node], X_shard)
  loss = \
    (IN_MEMORY_SHARD_SIZE / TRAINING_DATASET_SIZE) * mse(y_shard, y_est)
  loss.backward()

代码执行完毕后，

[W[node].grad for node in range(NODES)]

应该输出

[tensor([ -0.1776, -10.4762, -19.9037, -31.2003]),
 tensor([  0.0823, -10.3284, -20.6617, -30.2549]),
 tensor([ -0.1322, -10.9773, -20.4698, -30.2835]),
 tensor([  0.1597, -10.4902, -19.8841, -29.5041])]

代表环形集群中每个节点的模型梯度的张量。

请注意，在使用列表 8.7 中 for 循环中的代码对每个节点执行梯度下降的前向和后向步骤之后，Horovod 算法必须执行两个基于环形网络的阶段，以便将累积梯度通信到环中的每个节点。第一个阶段称为 reduce-scatter，在第 8.5 节中进行了解释，第二个阶段称为 all-gather，在第 8.6 节中进行了解释。

8.5 阶段 1：Reduce-scatter

本节介绍了 Horovod 的 reduce-scatter 阶段，假设环形集群中的每个节点都使用模型参数的相同副本进行初始化。本节继续使用列表 8.7 中的示例，其中模型参数的相同副本存储在 W[node] 中，并且每个节点完成了梯度下降的前向和后向步骤，导致的梯度值保存在 W[node].grad 中。通过本节的结束，您将了解 Horovod 的 reduce-scatter 阶段如何确保环中的每个节点最终获得累积梯度 g[0] + g[1] + g[2] + g[3] 的不同段。

Horovod 的第一个阶段称为 reduce-scatter，在每个节点都完成基于数据集的节点特定分片的梯度计算后开始。如前一节所述，每个节点逻辑上将计算的梯度分成 NODES 个段。此阶段的第一次迭代（共三次迭代）显示在图 8.6 中，其中图的顶部显示，在阶段开始时，每个节点 n[i] 存储着特定于分片的段，s0 到 s[NODES-1](n[i])。

图 8.6 第一次 reduce-scatter 阶段的迭代启动了跨节点的梯度段传输。

由于每次迭代 reduce-scatter 仅将一个段的数据发送到后续节点，因此在第一次迭代（如图 8.6 底部箭头所示）中，节点 n[i] 将段 s[(i - 1)] % NODES(n[i]) 转发给后续节点。在第一次迭代结束时（图 8.6 底部），每个节点 n[i] 累积了一个段 s(i - t - 1) % NODES % NODES]) + s(i - t - 1) % NODES，其中 t=1 代表第一次迭代。

在后续的迭代中，每个节点将上一次迭代中累积的段发送给后继节点。例如，在第二次迭代（如图 8.7 所示）中，节点n[1]发送段s[3]（n[0] + n[1]），节点n[2]发送段s[0]（n[1] + n[2]），一般来说，对于第 t 次迭代，节点n[i]发送累积的段s[(i - t)] % NODES(n[i])。由于在具有四个节点的示例中，只需要三次迭代来减少散布段，因此图 8.7 的底部显示，到第二次迭代结束时，每个节点上只缺少每个段的一部分：由s(i + 1) % NODES 指定的段。

图 8.7 第二次减少散布迭代传播累积梯度。

这个缺失的部分在示例的第三个和最后一个迭代中填补，即在迭代结束时（图 8.8 的底部），每个节点n[i]都累积了整个段s[i]。例如，注意在图 8.8 中，n[0]以s[0]结束了本阶段的最后一次迭代，节点n[1]以s[1]结束，依此类推。

图 8.8 第三次减少散布迭代完成四节点环的梯度累积。

列表 8.8 减少散布阶段的 Python 伪代码

for iter in range(NODES - 1):
  for node in range(NODES):
    seg = (node - iter - 1) % NODES     ❶
    grad = W[node].grad[seg]            ❷
    next_node = (node + 1) % NODES
    W[next_node].grad[seg] += grad      ❸

❶ 第一个段被累积到第一个节点。

❷ 检索与节点和段 seg 对应的梯度值。

❸ 在环中的下一个节点上累积梯度段的值。

在列表 8.8 中的代码执行完毕后，可以使用以下方式输出结果梯度

print([f"{W[node].grad}" for node in range(NODES)])

应该打印出

['tensor([ -0.0679, -31.9437, -39.7879, -31.2003])',
 'tensor([  0.0823, -42.2722, -60.4496, -61.4552])',
 'tensor([-4.9943e-02, -1.0977e+01, -8.0919e+01, -9.1739e+01])',
 'tensor([ 1.0978e-01, -2.1468e+01, -1.9884e+01, -1.2124e+02])'].

请注意，如预期的那样，梯度值分散在节点之间，以便n[0]存储累积梯度的段s[0]，n[1]存储段s[1]，依此类推。一般来说，减少散布后的梯度累积段可以使用以下方式打印出来

print([f"{W[node].grad[node]}" for node in range(NODES)]),

它在每个节点上输出累积段的值：

['-0.06785149872303009', '-42.27215576171875', '-80.91938018798828', '-121.24281311035156']

图 8.9 的插图总结了当减少散布环由四个节点组成时，列表 8.8 中的代码的情况。

图 8.9a 减少散布的迭代

图 8.9b 减少散布的迭代

大规模 MLOps 工程（三）（4）https://developer.aliyun.com/article/1517794