搜索
查看: 690|回复: 3

可编程数据平面加速分布式人工智能模型训练

[复制链接]
发表于 2024-4-12 17:34:22 | 显示全部楼层 |阅读模式
本帖最后由 中国计算机学会 于 2024-4-19 14:52 编辑

摘要—本文分析了分布式训练中由于网络通信造成的性能瓶颈,并阐述了如何利用可编程网络改善分布式训练的物理网络瓶颈问题。此外,从基础性的整体架构优化和细节性的部署容错优化两个方面,我们论述了目前已有方案对分布式训练的改进策略。最后,我们对未来的研究方向进行了展望。

吴春明浙江大学
栋(福州大学

关键词 :可编程数据平面 分布式训练 算网融合

  

引言

        人工智能大模型的高速发展对数据的实时处理与快速迭代提出了更高要求。可编程数据平面作为支撑人工智能大模型分布式训练的基础设施,为模型训练提供坚实保障,提升训练质量,有力地推动了自动驾驶网络等应用的发展与落地。
       作为算网融合中的重要组成部分,数据平面负责网络数据的存储与转发。随着人工智能技术的发展,传统软件定义网络中的数据平面已经不能很好地适配当前新的计算需求。由新型可编程设备如SmartNIC、可编程交换机、NetFPGA构成的可编程数据平面可提供更强大的计算能力与数据处理能力,如支持自定义数据包报头、对数据包特定字段进行哈希计算等操作。可编程网络优异的特性为加速机器学习分布式训练提供了可能。
        可编程网络的灵活控制特性让交换机可以承担简单的网络计算,卸载部分网络功能。例如,在神经网络模型分布式训练梯度迭代过程中,传统数据平面传输时延占据总时间开销的83.2%[1],极大地降低了模型的训练效率。为解决该问题,现有工作依托可编程数据平面实现参数收集、梯度聚合、梯度更新等分布式训练过程,通过优化服务器和工作节点之间的端到端通信时延,大幅提升模型的分布式训练效率。此外,可编程数据平面还可以基于不同神经网络的模型特征,制定定制化流量优化策略,在可容忍的精度范围内实现分布式训练加速。

可编程数据平面赋能人工智能模型训练

       近年来,神经网络模型的训练主要采用分布式计算架构,利用多台GPU服务器和物理网络协同实现分布式模型训练,打破单机训练的算力瓶颈。尽管如此,物理网络仍难免成为限制分布式训练的瓶颈,主要原因有两点。
       1.模型复杂性:当批处理大小(batch size)固定时,随着复杂神经网络层次和参数的增加,同步模型中worker(执行计算任务的工作单元)额外等待的时间也随之增长,导致批处理时间较长。同时,神经网络层数越高,批处理量越小。现有结果表明[2],复杂模型需要在更短的时间内传输更多的参数,增加了通信计算比,给物理网络造成了压力。
        2.额外的网络通信开销:在仅需少量worker就可进行模型训练的场景中,增加worker的数量会引入额外的网络通信时延,降低GPU的利用率,导致批处理的总时间增加。另外,由于分布式训练中的批处理引入了多机之间的通信,其批处理时间比单机训练时间更长,甚至会因为网络拥塞、数据包丢失等问题导致训练进程中断。因此无论当前使用的硬件配置和神经网络模型设置如何,网络都是一个显著的限制因素。
        纵观GPU和物理网络的发展历程,GPU的计算能力和内存容量都是逐年增长,但物理网络的发展速度相对缓慢,导致分布式训练中物理网络瓶颈问题愈加严峻。
       为突破这一瓶颈,研究者在可编程网络领域开拓了新的方向。利用可编程交换机灵活与可重构的特点,卸载分布式训练中网络通信部分功能,进而充分利用高性能网络资源,为分布式训练提供更好的解决方案。

       利用可编程网络优化分布式训练的方法主要有三种,如图1所示:(1)使用可编程交换机充当参数服务器(parameter server)执行网内聚合[3~5],优化分布式训练架构;(2)通过控制平面的决策优化分布式训练的部署和路由方案;(3)利用可编程交换机保障链路可靠性和管理梯度数据包的调度。下文阐述了近期基于可编程网络优化分布式训练的工作以及对未来工作的展望。
架构设计优化
        分布式训练架构的设计是整个加速过程的基础,为其他细节优化提供框架和主要方向,并成为优化方案中各部分之间的联系主体。架构的整体优化主要涉及三个方面:分布式网络结构、硬件结构设计和网络的可扩展性。

    分布式网络结构:从最新研究来看,分布式训练的网络结构主要有参数服务器[3, 4]和All-Reduce[6]两种。如图2所示,在参数服务器结构中,每个worker计算模型参数的梯度更新,并将这些更新发送给参数服务器。参数服务器通常是专用服务器,负责全局参数的聚合计算并将新参数值分发给所有worker。这种架构虽然易于扩展,但由于参数服务器扮演着全局参数聚合的中心角色,且受制于网络中存在的较多跳数,训练速度会受到限制。此外,参数服务器也容易成为整个聚合网络的瓶颈。

       All-Reduce方法是一种集体通信技术,如图3所示,所有的worker通过一个环拓扑进行通信交流。虽然随着worker数量的增加,这种方法会导致时延增加,但在带宽利用方面是最优的。

       在分布式架构与可编程数据平面结合方面,现有研究均是对参数服务器架构进行优化,如图4所示,可编程交换机充当参数服务器执行网内聚合加速分布式训练,例如iSwitch[1]、SwitchML[7]等工作都使用了这种新型网络结构。在该结构中,聚合操作被卸载到可编程交换机中。这种改进减少了交换机到参数服务器阶段的通信开销,并将梯度数据包的网络跳数减少了一半。因此,将分布式训练的延迟和通信成本最小化成为可能。
       ATP[8]在参数服务器网络结构的基础上,引入了网内聚合的方法。在梯度数据包传输的路径上,交换机以最大努力进行聚合,逐步减少了网络中的梯度流量,并且能够容忍一定程度的网络故障,极大地提升了分布式训练的网络性能。
      硬件结构设计:在传统的分布式训练中,交换机主要用于转发数据包,与一般网络中的交换机没有本质区别。然而,在引入网内聚合的分布式训练中,可编程交换机需要同时承担数据包转发和梯度数据包聚合的任务。为此,iSwitch对交换机硬件进行了修改,如图5所示。它增加了一个输入仲裁器,在交换机中判断输入数据包的类型,使交换机能够同时处理训练梯度包和普通数据包。此外,在聚合器的设计上,为了加快梯度数据的处理,并实现数据级并行性最大化,iSwitch从更细粒度的角度处理数据包(即在内部总线的单个时钟周期内可以操作的数据量)。因此,每个数据包都被划分为多个更小的单位以便于并行处理。



      网络可扩展性:随着分布式训练规模的不断扩大,计算节点的数量也随之增加。同时,面对多租户场景下的训练需求,训练网络应能依据需要进行动态伸缩,实现通过改变参数服务器的数量或构建参数服务器集群增减计算节点。然而,使用All-Reduce方法的网络在可伸缩性方面表现较差。它能够适应训练规模的缩小,但随着训练规模的增大,其网络性能将受到较大的影响。因此,为了在网内聚合训练中满足这种动态需求,iSwitch提出了分层聚合机制,将所有的交换机分为三层,自下而上分别是:单机架下的ToR(Top of Rack,机架顶部)交换机、聚合层交换机、核心层交换机,如图6所示。同一机架上的所有服务器通过ToR交换机进行连接;聚合层交换机负责梯度包路径上的聚合工作,ToR交换机和聚合层交换机在完成缓冲区梯度向量的本地聚合后,将聚合后的向量段发送至更高一层交换机,即核心层;核心层交换机负责最终的全局聚合以及最终更新参数数据包的广播。这种分层机制更易于分布式训练规模的伸缩,能够动态适应大规模多租户场景下的训练。



部署与容错优化
       分布式架构设计的优化为训练过程中其他方面的优化提供了坚实的基础。然而,真实训练场景需要更细致的部署和容错优化设计,以保障整体架构设计方案的稳定性与可靠性,使分布式训练过程的加速真正可行。
        聚合部署:近期研究中,多数解决方案主要通过固定默认路由实现网内聚合。梯度流量的动态性和固定路由下传输链路的不可变性会导致训练过程负载不均,增加节点间的通信时长。为了适应不同类型任务的特点,GRID[9]考虑了训练过程中梯度数据包的路由设计,控制平面利用整数线性规划(Integer Linear Program,ILP)求解决策,确定梯度数据包的聚合路由。其中,我们以worker的最小发送速率最大化为目标函数,并且考虑以下五个条件约束。(1)聚合约束:假设每个梯度仅在网络中聚合一次,以平衡聚合效果和路径带宽消耗;(2)发送速率约束:所有worker的发送速率不得超出限制的最大发送速率;(3)处理能力约束:每个交换机在有限的处理速率下聚合梯度;(4)聚合节点约束:worker发送速率不得超过其对应聚合节点的发送速率;(5)带宽约束:转发速率不得超过接收节点的入口带宽。最终根据上述目标及约束进行ILP求解,得出最佳部署方案,下发到数据平面并完成部署。


      流量调度:在大规模的分布式训练中,梯度包流量呈现不断增长的趋势。为了降低网络性能在训练过程中受到超大流量的影响,并考虑到不同租户和任务的紧急程度差异,研究人员提出了一种名为Geryon[10]的解决方案。如图7所示,该方案利用多个具有不同优先级的流传输紧急程度不同的梯度数据包,并在整个网络结构中根据优先级处理这些数据包。具体而言,Geryon为每个梯度数据包设置优先级,并通过相应的优先级调度机制确保高优先级数据包能够优先转发,从而加快训练过程。考虑到交换机的处理资源有限,它会根据紧急程度在阈值上下把优先级分为高、低两个级别,由此降低交换机的调度开销。Geryon还考虑了神经网络模型中不同结构参数的紧急程度差异,以设置梯度数据包的优先级。同时,为了适应不同模型的需求,Geryon利用单轮迭代时间评估紧急程度阈值的合理性,并在初始几轮迭代中选择最优阈值,以供后续训练过程使用。最后,研究人员还希望该部署方案拓展至用户级别,根据用户需求设置任务优先级,将宏观和微观的优先级调度相结合。
      丢包处理与拥塞控制:在现实网络中,网络流量或设备故障等因素通常会导致网络拥塞、数据包丢失等情况发生,这会极大地损害分布式训练的性能,甚至使分布式训练系统停止运行。因此,分布式训练系统中的丢包处理与拥塞控制极为重要。
       针对丢包处理问题,SwitchML考虑到交换机资源的限制,在主机端引入了超时重传机制。然而,简单的主机端重传机制会导致交换机产生不必要的重复聚合或干扰其他梯度片段的聚合。为此,SwitchML首先区分处于上传聚合和广播更新阶段的丢包,通过在交换机记录已被聚合的数据包防止重复聚合。其次,为每个在交换机端完成聚合的梯度数据维护一个副本,以用于广播更新阶段的丢包重传处理。这样一来,SwitchML有效地解决了丢包问题,并保证了分布式训练的顺利进行。
       在拥塞控制方面,针对大规模多租户场景下各类网络资源严重争抢的问题,交换机中聚合器的高度竞争可能导致无法完全聚合所有流量,从而出现交换机队列长度的累积和缓冲区溢出引发的数据包丢失。为了解决这一问题,ATP采用显式拥塞通知(Explicit Congestion Notification,ECN)作为拥塞信号,并将其插入梯度数据包中。完成聚合后,将ECN标志复制到用于参数更新的广播数据包中,使所有节点能够记录ECN表示的拥塞状态。每个worker根据接收到的拥塞信息,利用线性增加/倍数减少的方法调整发送窗口大小,以应对拥塞情况。通过这种方式,系统能够有效地实现拥塞控制,以提高网络性能和降低数据包丢失的风险。

未来工作

        随着人工智能模型规模的不断扩大,分布式训练需要更多的计算资源和网络节点,这对可编程数据平面提出了更高的要求。为了更好地满足人工智能模型分布式训练的需求,实现分布式训练中网络通信与模型训练协同一致,我们从新型可编程数据面构建、网内聚合架构优化、模型与网络的高协同性三个方面详细阐述对未来分布式训练网络的展望。
      算力网络中的新型数据面构建。算力感知网络作为面向未来的新型网络架构,实现了网络和计算的相互感知、高度协同。设计基于SRv6的算力路由协议,构建基于IPv6/SRv6等多种协议的可编程数据平面,实现支持网络编程、灵活可拓展的新型数据面,是算力感知网络未来的重要研究方向。
      多租户场景下的网内聚合架构优化。现有方法大多针对单任务进行训练加速。当任务拓展至大规模的多租户任务训练时,当前的架构设计就难以支持高效并发训练。为了适应分布式训练发展趋势,保证未来应用中分布式训练的稳定运行和可靠性,须从整体架构入手,重新考虑各个优化步骤,使其更适配未来的分布式训练需求。
      分布式训练中的高协同性网络构建。在分布式训练中,网络特性与数据中心存在差异。因此,在构建分布式训练网络时,应充分考虑其结构特征、流量模式和传输需求的因素。目前分布式训练网络具有以下几个特点:流量模式具有可预测性和模型依赖性;参数计算在精度上具有较高容错性;梯度数据包具有顺序无关性;不同阶段的数据包对训练产生的影响程度和紧急程度不同。因此,未来的分布式训练网络设计应从其独特性出发,尽可能与训练需求契合,以实现更好的协同加速效果。  ■
参考文献:
[1] Li Y, Liu I J, Yuan Y, et al. Accelerating distributed reinforcement learning with in-switch computing[C]// Proceedings of the 46th International Symposium on Computer Architecture (ISCA 2019). 2019: 279-291.
[2] Luo L, Liu M, Nelson J, et al. Motivating in-network aggregation for distributed deep neural network training[C]// Workshop on Approximate Computing Across the Stack. 2017.
[3] Li M, Zhou L, Yang Z, et al. Parameter server for distributed machine learning[C]// Big Learning NIPS Workshop. 2013, 6(2).
[4] Luo L, Nelson J, Ceze L, et al. Parameter hub: a rack-scale parameter server for distributed deep neural network training[C]// Proceedings of the ACM Symposium on Cloud Computing (SoCC 2018). 2018: 41-54.
[5] Dean J, Corrado G, Monga R, et al. Large scale distributed deep networks[C]// Proceedings of the 25th International Conference on Neural Information Processing Systems. 2012: 1223-1231.
[6] Li Y, Park J, Alian M, et al. A network-centric hardware/algorithm co-design to accelerate distributed training of deep neural networks[C]// Proceedings of the 51st Annual IEEE/ACM International Symposium on Microarchitecture (MICRO 2018). 2018: 175-188.
[7] Sapio A, Canini M, Ho C Y, et al. Scaling distributed machine learning with in-network aggregation[C]// Proceedings of the 18th USENIX Symposium on Networked Systems Design and Implementation (NSDI 21). 2021: 785-808.
[8] Lao C L, Le Y, Mahajan K, et al. ATP: In-network aggregation for multi-tenant learning[C]// Proceedings of the 18th USENIX Symposium on Networked Systems Design and Implementation (NSDI 2021). 2021: 741-761.
[9] Fang J, Zhao G, Xu H, et al. GRID: Gradient routing with in-network aggregation for distributed training[J]. IEEE/ACM Transactions on Networking, 2023: 2267-2280.
[10] Wang S, Li D, Geng J. Geryon: Accelerating distributed CNN training by network-level flow scheduling[C]// IEEE Conference on Computer Communications (INFOCOM 2020). 2020: 1678-1687.


版权声明:中国计算机学会(CCF)拥有《中国计算机学会通讯》(CCCF)所刊登内容的所有版权,未经CCF允许,不得转载本刊文字及照片,否则被视为侵权。对于侵权行为,CCF将追究其法律责任。






回复

使用道具 举报

0

主题

2

回帖

0

好友

管理员

发表于 2024-4-13 18:25:27 来自手机 | 显示全部楼层
nice
回复

使用道具 举报

1

主题

2

回帖

0

好友

新手上路

发表于 2024-4-13 19:04:41 来自手机 | 显示全部楼层
好棒,学习了
回复

使用道具 举报

0

主题

2

回帖

0

好友

管理员

发表于 2024-4-29 11:34:16 | 显示全部楼层
very good
回复

使用道具 举报

您需要登录后才可以回帖 登录

手机版|CCF Link ( 版权所有 中国计算机学会  京ICP备13000930号-4|京公网安备 11010802032778号   )

GMT+8, 2025-4-27 07:07 , Processed in 0.047879 second(s), 20 queries .

快速回复 返回顶部 返回列表