蚂蚁提出自适应训练框架AntDT，有效解决工业级分布式训练快慢机的利器

2024-04-24 06:04

近⽇，蚂蚁集团语言与机器智能部-知识引擎团队的论文《AntDT: A Self-Adaptive Distributed Training Framework for Leader and Straggler Nodes》中稿 ICDE 2024 (Industrial Track)。

ICDE (IEEE International Conference on Data Engineering) 是电气与电子工程师协会（IEEE）举办的旗舰会议，与 SIGMOD、VLDB 并称数据管理与数据库领域的三大国际顶级学术会议，入选为中国计算机学会推荐的 A 类国际学术会议，在国际上享有盛誉并具有广泛的学术影响力。

蚂蚁集团语言与机器智能部-知识引擎团队，长期聚焦于金融场景特色 AI 基础引擎的建设，包括图学习、因果学习、知识图谱、AI 对齐等相关的基础框架和引擎。团队目前聚焦于大模型和知识图谱双向驱动框架和系统的研发。

论文链接：

https://arxiv.org/abs/2404.09679

随着深度学习在越来越大的数据上取得巨大成功，分布式训练技术越来越受到工业界的关注，尤其是数据并行架构，用来加速训练过程。然而大规模训练中，训练的效率可能会受到 straggler 问题（或快慢机问题）的显著影响，造成慢节点拖累正常节点训练，严重影响训练效率。

围绕该课题，蚂蚁集团语言与机器智能部推出论文《AntDT: A Self-Adaptive Distributed Training Framework for Leader and Straggler Nodes》，该工作提供了一个自适应的训练框架 AntDT，以系统地解决工业级分布式训练中不同类型的慢节点问题，同时隐藏了由不同慢节点缓解方法引起的数据分配和容错的复杂细节。

此外，基于 AntDT 框架，我们提出了两种慢节点缓解解决方案作为运行示例，在 worker 节点和 server 节点上有效地缓解慢节点问题，助力高效的工业级分布式训练。

背景

在分布式训练中，现有数据并行有两种主要架构，参数服务器 (Parameter Server/PS) 和 AllReduce 架构。然而，在大规模分布式训练中，训练的效率可能会受到 straggler 问题（或快慢机问题）的显著影响。Straggler 节点指的是慢节点，其阻碍了其他正常节点（也称为 Leader 节点）的进展，直到它们完成其分配的 workloads。

这个问题出现的主要原因有两个，即硬件异构性和资源争用。硬件异构是第一个原因，特别是对于拥有较旧和较新系列设备组成的集群而言，我们将这些慢节点称为确定性慢节点。资源争用是导致慢节点的另一个原因，这会引起非确定性的慢节点。

云供应商主要使用两种集群，分别是非专用集群和专用集群。非专用集群是指被多个租户共享的集群，例如混部集群或者抢占式实例(spot instance)，而专用集群由专用设备组成，专门为一个租户独占。这些慢节点是非确定性的，我们根据慢节点的时间周期性将它们分类为短期慢节点和长期慢节点。

作为一个例子，我们首先将参数服务器中的快慢机问题进行形式化。假设给定全局 Batch size 为 B，n 个 worker 节点中的一个节点 i 首先根据本地 batch size 大小的数据（Bi）。

我们将 worker 节点 i 在一个迭代中的整个时间称为批数据处理时间（BPT）。它可以分解为三部分：worker 节点上的计算时间（包括基于批量数据 Bi 的前向和反向传播时间），其可以表示为批量数据 Bi 的函数 F（Bi）；worker 节点 i 在 server 节点上的计算时间（即 server 节点聚合来自 worker 节点 i 的局部梯度并更新模型参数）；以及 worker 节点 i 和 server 节点之间的通信时间，即测量 worker 节点拉取最新参数并推送局部梯度进行同步的通信时间。

那在 BSP 训练模式中，整个训练过程可以形式化为以下优化问题。给定 n 个 worker 节点和 m 个 server 节点，我们希望最小化所有 worker 节点中最长的批数据处理时间 BPT，即：

要完全解决快慢机问题，我们面临三个主要挑战。首先，在公式 1 中，需要同时最小化，和这三项，而现有的工作只能最小化这些项的一部分，未能完全消除 stragglers。

首先，为了减少，基于负载均衡的方法（如 LB-BSP）尝试重新平衡分配给慢节点和快节点的工作负载。这通常具有较低的时间成本，并且对于短期的慢节点和确定性的慢节点都是有效的。同样，基于复制的方法（如 backup workers）通常启动已经被识别为慢节点任务的副本任务，并且只接受作业中第一个完成的任务的结果。但该方法会在废弃的慢节点中丢失一些样本，因此损害模型的统计效果。

此外，这些方法需要一种数据分配机制，可以有效地将样本分配给不同的 worker 节点，因为在快慢机之间的样本的消耗不平衡。然而，这些方法依赖于定制和复杂的数据分配机制，使它们彼此不兼容。此外，这两种方法对于由计算或网络资源争用引起的长期慢节点是无效的，也无法有效解决和。

因此，从集群调度器的角度来看，基于调度的方法（如 kill-restart）尝试通过杀死慢节点，并重新启动新的节点（这需要从定期保存的 checkpoints 中恢复训练），这可以在资源争用时，减少，和这几项。然而，这种方法非常耗时，因为它包括调度和恢复训练状态的时间。因此，它对于短期慢节点会产生很高的时间成本。

其次，一个自适应分配数据的机制对于确保各种解决快慢机问题方法的可扩展性和兼容性至关重要。这些相关方法采用复杂的数据分配策略，主动地将数据负载从慢 worker 节点转移到更快的 worker 节点，这些策略彼此不兼容，在工业环境中效率低下。因此，该框架需要一个灵活的数据分配机制来利用所有这些方法。

其三，基于复制和调度的方法在训练过程中可能会丢失样本，从而损害了由“至少一次”和“至多一次”语义构成的数据完整性。特别地，“至少一次”意味着每个样本至少应被使用一次，而“至多一次”意味着在训练过程中对于一个周期（整个数据集的一次完整迭代）来说，每个样本不应该被使用超过一次。特别是，在金融应用中的机器学习任务中，“至少一次”语义比“至多一次”语义更为重要。

尽管多消费数据可以被视为机器学习应用中的过采样，但对慢 worker 节点的数据进行丢弃在许多场景下是不可接受的；否则，它会损害数据集的完整性，并可能影响统计效果。例如，欺诈检测严重依赖于机器学习技术，其中正样本（欺诈交易）明显少于负样本（正常交易），而丢失这样的正样本被认为是不可接受的，特别是在金融应用中。

方案

a）AntDT框架

1. 我们首先概述整个框架的设计。该框架由四个组件组成：带状态的数据分片服务（Stateful DDS）、AntDT Monitor、Controller 和 Agent，如图所示。它提供了数据分配和容错等基本功能，具有高效性和可扩展性，并与具体的慢节点处理操作解耦。基于我们的框架，用户可以进一步利用预定义的慢节点处理方案，或根据实际情况轻松自定义特定的慢节点处理方案。框架部分包含：

Stateful Dynamic Data Sharding Service (Stateful DDS): 动态数据分片的主要思想是使用动态数据分区策略，以细粒度地将相应的数据分片分配给计算能力不平衡的 worker 节点。此外，每个数据分片都拥有一个状态，以确保在任何节点终止的情况下保持数据完整性。然后，我们可以在分片级别管理训练数据，包括灵活的数据分配和数据完整性。

2. AntDT Monitor: AntDT Monitor 组件旨在从不同的数据源收集可观察信息，以供 Controller 进一步进行慢节点处理。AntDT Monitor 将定期收集和聚合三种类型的信息用于慢节点检测，包括应用程序状态（训练吞吐）、节点状态和来自其他模块或第三方 API 的信息。

3. AntDT Controller: 我们在 AntDT Controller 中提供了预定义的慢节点处理操作集。这包括节点操作，例如 KILL RESTART，和全局操作，例如 ADJUST BS（调节 batch size）、BACKUP WORKERS（增加 backup workers）和ADJUST LR（调学习率），这些操作需要在节点之间进行同步。

4. AntDT Agent: AntDT Agent 旨在为 Monitor 收集信息并异步执行 Controller 发送的操作。它部署在每个 worker 节点或 server 节点上。一方面，它异步收集和推送相应的应用程序和节点信息给 Monitor。另一方面，Agent 执行 Controller 通知的慢节点处理操作。在我们的实现中，Agent 进程通过环境变量与 worker 节点或 server 节点的训练进程进行通信。

b）解决方案

基于 AntDT 框架，我们设计和实现了两种慢节点处理方案。这两个方案，作为运行示例，展示了如何根据蚂蚁集团的集群状态与 AntDT 框架协同工作来解决慢节点问题。这些解决方案在实践中可以很容易地扩展到其他集群。

第一个解决方案是 AntDT-ND 方案（用于非专用集群），其利用非专用集群中慢节点的时间周期性，采取轻量级的 ADJUST BS 操作来减轻短期慢节点，并谨慎地采取重量级的 KILL RESTART 操作来消除长期慢节点，提高训练效率。它也同时解决了参数服务器在 worker 节点和 server 节点上的慢节点问题。

第二个解决方案是，我们提出了 AntDT-DD 方案（用于专用集群），进一步通过梯度累加和混合整数规划，优化了 ADJUST BS 操作，以在专用 GPU 集群中，最大化 AllReduce 架构的训练速度。

实验效果

a）训练效率：

i. worker 慢节点：

在 BSP 训练中，原生 TensorflowBSP 受到 worker 节点异常的严重影响，训练时间达到 8144 秒（超过 2 小时）。如图 10 所示，与 Backup Worker 方法、LB-BSP 和 BSP 相比，AntDT-ND 在 BSP 训练中能够平均缩短总体训练时间 24%、44% 和 51%。

在 ASP 训练中，AntDT-ND 表现优于 ASP-DDS 和 ASP，分别提高了 16% 和 325%，如图 11 所示。相比于 LB-BSP 仅调节 batch size，其主要贡献来自 DDS 服务通过动态数据分片在 worker 节点之间自适应调整工作负载。其次的贡献是通过执行 KILL RESTART 操作来替换长期慢节点。

ii. server 慢节点：

AntDT-ND 在 BSP 训练中比所有其他方法快两倍以上。具体而言，AntDT-ND 将 LB-BSP、Backup Worker 方法和 BSP 的 JCT 分别减少了 51%、66% 和 61%，如图 10 所示。在 ASP 训练中，我们观察到与图 11 中显示的类似趋势，AntDT-ND 的速度是原生 ASP 的两倍。

此外，我们注意到，当只有 server 节点慢节点时，ASP 甚至比 BSP 更耗时，这与直觉相反。这是因为 ASP 要求在 server 节点端更频繁地更新模型参数，而任何 worker 节点在完成本地计算后都将与所有 server 节点通信。

b）鲁棒性

我们进一步评估 AntDT-ND 在训练效率方面的鲁棒性，通过分析其在不断增加的慢节点强度下的表现。首先，我们将 worker 节点的慢节点强度从 0.1 提升到 0.8。表 III 显示，原生 BSP 的 JCT 上升，而 AntDT-ND 的 JCT 仅略微增加。与 BSP 相比，当我们将 worker 节点的慢节点强度从 0.1 提升到 0.8 时，AntDT-ND 的速度提高了 10.3% 至 104.5%。此外，AntDT-ND 的 JCT 变化相对于 BSP 来说较小。

c）Evaluation of AntDT-DD

我们在混合的 GPU 设备上训练模型，以评估 AntDT-DD 在异构专用集群中的工作情况。如图 15 所示，DDP 花费大约 1266 秒，LB-BSP 花费约 1020 秒，而 AntDT-DD 只需 912 秒就能完成使用 ResNet-101 的 ImageNet 的一个 epoch 的训练。因此，相较于相同数量的训练数据，AntDT-DD 在训练速度上超过了 DDP 和 LB-BSP，分别提高了 38.8% 和 12%。

对于通信密集型的模型，如 MobileNets，性能差距扩大。AntDT-DD 在 MobileNets 中的运行速度比 LB-BSP 和 DDP 快 25% 和 48.5%。总而言之，与异构 GPU 集群中的其他方法相比，AntDT-DD 可以实现近 1.5 倍的加速。

d）Overhead Analysis and Scalability

我们进一步评估了 AntDT 框架在扩展到数百个节点时的开销。AntDT 的开销主要来自于 Stateful DDS 服务和 AntDT Agent 中的同步时间。我们通过总同步时间除以过程中的 JCT 来衡量额外开销（不考虑异步线程的时间消化）。我们以 BSP 训练中 AntDT-ND 的延迟时间为例，因为 AntDT-DD 中的 batch size 调整过程仅运行一次，开销可以忽略不计。

我们分别报告了三个非专用生产 CPU 集群规模中的开销百分比。在小规模集群中，延迟时间占 JCT 的 0.46%，其中 55% 是状态同步在 Stateful DDS 中，45% 是同步机制的时间。在中等规模集群和大规模集群中，开销百分比略有波动，并且都小于总体 JCT 的 0.44%。总之，实验结果验证了 AntDT 框架的开销几乎可以忽略不计，并且可以在大规模分布式训练中进行扩展。

总结

为了系统地解决实际工业场景中分布式训练的快慢机问题，我们提出了一个自适应的训练框架 AntDT，其可以轻松地使用不同的慢节点缓解方法来定制解决方案，而无需考虑复杂的数据分配和容错机制。其次，本文提出了两种解决方案运行示例，其使用我们的框架解决了实际生产集群中各类型的慢节点，这两个解决方案在对比其他 SOTA 方法时提供了 3 倍以上的性能提升。我们广泛的实验结果清楚地证明了该框架在实际工业场景中的有效性。

更多阅读