蚂蚁提出自适应训练框架AntDT,有效解决工业级分布式训练快慢机的利器
近⽇,蚂蚁集团语言与机器智能部-知识引擎团队的论文《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 节点上有效地缓解慢节点问题,助力高效的工业级分布式训练。
背景
其次,一个自适应分配数据的机制对于确保各种解决快慢机问题方法的可扩展性和兼容性至关重要。这些相关方法采用复杂的数据分配策略,主动地将数据负载从慢 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 架构的训练速度。
实验效果
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 节点通信。
我们进一步评估 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,其可以轻松地使用不同的慢节点缓解方法来定制解决方案,而无需考虑复杂的数据分配和容错机制。其次,本文提出了两种解决方案运行示例,其使用我们的框架解决了实际生产集群中各类型的慢节点,这两个解决方案在对比其他 SOTA 方法时提供了 3 倍以上的性能提升。我们广泛的实验结果清楚地证明了该框架在实际工业场景中的有效性。
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。
总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。
📝 稿件基本要求:
• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注
• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题
• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道:
• 投稿邮箱:[email protected]
• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿
△长按添加PaperWeekly小编
🔍
现在,在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
微信扫码关注该文公众号作者