翻译自八卦文档

八卦已经整合的原语包括

  • 集中式同步通信(AllReduce)
  • 去中心化同步通讯
  • 低精度通信

其有效性已经在各种场景和模型中得到验证,包括ImageNet上的VGG和ResNet,Bert Large,以及快手等多个大规模工业应用:

  • 支持数十个TB参数模型训练的推荐系统,
  • 超过 10 亿个图像/视频的视频/图像理解,
  • ASR 与 TB 级数据集,
  • 等等。

使用举例

使用bagua类似于使用其他分布式训练库,如 PyTorch DDP 和 Horovod。

def main():
    args = parse_args()
    # define your model and optimizer
    model = MyNet().to(args.device)
    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr)
    # transform to Bagua wrapper
    from bagua.torch_api.algorithms import gradient_allreduce
    model = model.with_bagua(
        [optimizer], gradient_allreduce.GradientAllReduceAlgorithm()
    )

    # train the model over the dataset
    for epoch in range(args.epochs):
        for b_idx, (inputs, targets) in enumerate(train_loader):
            outputs = model(inputs)
            loss = torch.nn.CrossEntropyLoss(outputs, targets)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

计算使用的是pytorch, 只需要包上一个通信算法原语就可以了。

通信算法

八卦依靠分布式学习算法的多样性而蓬勃发展。 系统的极大灵活性使得可以平滑地结合各种SOTA算法,同时在执行过程中提供性能的自动优化。 对于最终用户,八卦提供了广泛的算法选择,她可以轻松地尝试执行她的任务。 对于算法开发人员来说,八卦是一个游乐场,在那里她可以只专注于算法本身(例如,逻辑和控制),而无需在不同的算法之间重新发明轮子(例如,通信原语和系统优化)。

梯度 allreduce

总览

Gradient AllReduce 算法是一种流行的同步数据并行分布式算法。 它是大多数现有解决方案中实现的算法,比如 PyTorch DistributedDataParallel, Horovod, 以及TensorFlow Mirrored Strategy

算法

使用此算法,每个worker在每次迭代中执行以下步骤。

  • 使用小批量计算梯度。
  • 使用 AllReduce 集合计算所有worker的梯度平均值。
  • 用平均梯度更新模型。

在八卦中,这个算法是通过 GradientAllReduce 算法类来支持的。 八卦中 GradientAllReduce 实现的默认性能应该与 PyTorch DDP 相当,并且在大多数情况下比 Horovod 更快。 八卦支持额外的优化,例如可以在实例化 GradientAllReduce 类时配置的分层通信。 在某些情况下,例如当机器间网络成为瓶颈时,它们可以使八卦比其他实现更快。

举例

运行 Gradient AllReduce 的完整示例可以在带有 --algorithm gradient_allreduce 命令行参数的 八卦示例中找到。

您需要初始化八卦算法(请参阅 API 文档了解您可以自定义哪些参数):

from bagua.torch_api.algorithms import gradient_allreduce
algorithm = gradient_allreduce.GradientAllReduceAlgorithm()

然后用以下方法装饰您的模型:

model = model.with_bagua([optimizer], algorithm)

bytegrad

总览

大规模分布式训练需要大量的通信成本,对于大型模型尤其如此。 例如,在使用 AllReduce 同步梯度的传统同步分布式设置中(许多库都是这种情况,比如 HorovodPyTorch DDP), 在训练过程的每次迭代中,需要在每个 worker 上发送和接收大小等于模型大小的梯度。 这样的通信成本很快成为很多场景下的训练瓶颈。

有很多关于如何应用模型/梯度压缩来节省这种通信成本的现有论文。 八卦提供了一种内置的梯度压缩算法,称为 ByteGrad,它在通信之前将梯度浮点数压缩到 8bit 字节。 这节省了原始成本的 3/4。 它通过优化的 CUDA 内核和分层通信实现了高精度的 min-max 量化算子。 这使得它比现有框架中(比如 PyTorch PowerSGD)的其他压缩实现快得多(在我们的基准测试中大约快 50%)。并且在相同的时期数 ByteGrad 在大多数任务上收敛近似于全精度算法。

算法

ByteGrad 在每次迭代中执行以下步骤。 假设我们有 m 个节点,每个节点有 n 个 GPU。

  1. 为所有i,ji,j 计算第 i 个节点的第 j 个 GPU 上的梯度 gi,jg_{i,j}
  2. 每个节点上的第一个 GPU 做一个 reduce 操作来计算同一节点上所有 GPU 的梯度的平均值,定义为第 i 个节点的 GiG_i
  3. 第 i 个节点上的第一个 GPU 使用量化函数 Q()Q(Gi)Q(⋅):Q(G_i ) 对所有 i 量化梯度 GiG_i。 然后每个节点在节点之间交换量化版本,使得每个节点都有所有 Q(Gi)Q(G_i ) 的平均值
  4. 每个节点上的第一个 GPU 将所有 Q(Gi)Q(G_i) 的平均值广播给同一节点上的所有其他 GPU,并且所有 worker 上的所有 GPU 都使用此量化平均值来更新模型

量化函数 Q(⋅) 计算其输入的最小值 x 和最大值 y,并将 [x,y] 拆分为均匀间隔的 256 个区间。 然后用一个 8 位整数表示其输入的每个元素,表示原始元素在哪个区间。

用法举例

您需要初始化八卦算法(请参阅 API 文档了解您可以自定义哪些参数):

from bagua.torch_api.algorithms import bytegrad
algorithm = bytegrad.ByteGradAlgorithm()
model = model.with_bagua([optimizer], algorithm)

去中心化SGD

去中心化训练概述

Decentralized SGD 是一种数据并行的分布式学习算法,它消除了所有worker之间集中式全局模型的要求,这使得它在通信模式上与基于 Allreduce 或基于参数服务器的算法有很大不同。 使用去中心化 SGD,每个 worker 只需要与一个或几个特定的 worker 交换数据,而不是全局聚合数据。 因此,去中心化通信的通信连接数比 Allreduce 少得多,通信开销比 Parameter Server 更均衡。 尽管去中心化 SGD 可能会导致每个worker的模型不同,但理论上已经证明,去中心化 SGD 算法的收敛速度与其中心化对应版本相同。 您可以在我们的 论文中找到关于去中心化 SGD 的详细分析 。

去中心化训练算法

目前,不时有许多分散的训练算法被提出。 这些令人惊叹的作品集中在分散训练的不同方面,如peer选择、数据压缩、异步等,并提供了许多有前景的见解。 到目前为止,八卦已经整合了两种基本的去中心化算法,即去中心化 SGD 和低精度去中心化 SGD。 凭借八卦对去中心化的自动系统支持,我们预计在不久的将来会实现越来越多的去中心化算法。

Decentralized SGD

现在我们将描述在八卦中实现的去中心化 SGD 算法。 假设worker数量为n,worker i上的模型参数为x(i),i0,...,n1x^{(i)},i∈{0,...,n−1}。 每个worker 都能够直接从任何其他worker发送或接收数据。 在每次迭代 t 中,算法重复以下步骤:

  1. 每个worker i 计算迭代的局部梯度 gt(i)g_t^{(i)}.
  2. 将本地模型与其选定的同伴的模型进行平均(表示为 xt(j)x_t^{(j)}), xt+12(i)=t(i)+xt(j)2x_{\frac{t+1}{2}}^{(i)}=\frac{t(i)​+xt(j)}{2}.
  3. 使用局部梯度更新平均模型.xt+1(i)=xt+12(i)γgt(i)x_{t+1}^{(i)}=x_{\frac{t+1}{2}}^{(i)}-\gamma g_t^{(i)}

在步骤 2 中,我们采用一种策略为每次迭代中的每个 worker 选择一个 peer,这样所有 worker 都正确配对并且数据交换是有效的,因为每个 worker 可以在迭代之间与不同的 peer 交换数据。 简而言之,我们的策略将worker 平均分成两组,并在两组之间动态配对工worker,每次迭代都不同。

点评,最后以谁的参数作为最终的ckpt呢?

通信开销

去中心化 SGD 的通信开销与网络程度高度相关,即一个worker与其他worker的连接数。 不同的拓扑或策略会导致不同程度的网络。 很明显,我们之前描述的Decentralized SGD算法的网络度为1。因此,在每次迭代中,一个worker只需要与一个worker建立一个连接来交换模型大小的1倍。 我们比较了不同通信模式在最繁忙节点的延迟和带宽方面的通信复杂性。

算法 时延复杂度 带宽复杂度
Allreduce (Ring) O(n) O(1)
ps O(1) O(n)
八卦去中心化SGD O(1) O(1)

用法举例

您需要初始化八卦算法(请参阅 API 文档了解您可以自定义哪些参数):

from bagua.torch_api.algorithms import decentralized algorithm = decentralized.DecentralizedAlgorithm()
model = model.with_bagua([optimizer], algorithm)

其他的算法

八卦提供的其他的算法没什么好说的,可以参考原文

八卦还支持弹性训练, 看上去是保存中间状态实现的。