在网聚合集群中多机器学习作业的管理策略

中国计算机学会

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

  
吴文斐（北京大学）

关键词 ：在网聚合 机器学习 作业管理策略

  

背景

        随着网络设备的可编程性和计算能力的提升以及构建分布式机器学习训练系统需求的出现，学术界提出一种新的计算传输原语——在网聚合（In-Network Aggregation，INA）。在网聚合完成的集合通信中的AllReduce和Reduce操作，是分布式机器学习模型同步过程中的必要原语。在在网聚合系统中，当多个张量以分组（packet）流的形式经过相同交换机时，交换机能够将张量聚合为一个结果（例如求和），并将结果发给接收者（Reduce）或者返回给所有发送者（AllReduce）。这个过程需要占用交换机内存资源，而一个作业的执行效率也受其能获取的内存大小影响。

        在生产环境中，为了有效利用有限的集群资源，使用者往往会同时运行多个作业。而在有在网聚合的集群中，交换机内存资源向多作业的分配，对每个作业的效率和所有作业的整体效率起关键作用。例如，交换机内存资源为0时，ATP（面向多租户的深度学习训练聚合传输协议）会回退到与无在网聚合集群相同的模式，而SwitchML（Barefoot公司主导提出的交换机卸载系统）则无法工作[1]。

        因此，在网聚合向生产集群部署的过程中，需要设计面向在网聚合作业的资源管理系统。这套系统能够实施不同的资源份内策略，提升整体作业效率。

在网聚合方案分类和对比

       目前，研究人员提出了多种在网聚合方案，它们对交换机资源的使用方法可以分为两类：独享模式和共享模式（见图1）。二者有共同的数据结构，但是工作流程不同。

        在网聚合方案中，多个发送者要发送相同长度的向量，每个发送者将向量切分为一个分组序列，一个序列中的分组按照逐个递增1的方式被设置一个分组序列号（Packet Sequence Number，PSN），所有的序列有相同的初始序列号。网络交换机的内存被组织为一个数组，每个单元称为聚合器，每个聚合器的存储空间与分组的载荷大小相同。发送者维护一个发送窗口将分组发送到交换机，完成聚合；聚合完成后，聚合结果分组被发送到接收者或者广播给发送者。发送者收到确认字符（ACKnowledge character，ACK）或者结果后，继续滑动窗口并发送新的分组。

独享模式

        如图1（a）所示，在该模式下，每个作业独占一片交换机内存，即聚合器数组的一个分区；分区之间互不重叠。每个分组按照取模的方式被映射到一个聚合器，在一个作业的区间(offset, size)中，映射关系为：aggregator.index = pkt.PSN % Size + Offset。

       一个作业中来自不同发送者的、PSN相同的分组被映射到相同的聚合器，并完成逐个分组的聚合；聚合完成后，结果被发送到下一跳设备。这种模式意味者发送端窗口不能超过数组大小；而为了正确处理丢包导致的发送者窗口进度不一致（progress differently），需要保证窗口大小不超过数组大小的一半。SwitchML[2]、NetReduce[3]等方案即采用此种模式。

共享模式

        如图1（b）所示，在该模式下，所有的作业共享整个聚合器数组，在整个数组空间内随机寻址。每个分组通过一个哈希函数映射到数组中的一个单元，映射关系为：aggregator.index = Hash(pkt.JobID, pkt.PSN)。一个作业中来自不同发送者的、PSN相同的分组仍然被映射到相同的聚合器。但是，不同作业不同PSN分组可能发生哈希冲突，对此，分组采用先到先服务的方式抢占聚合器，聚合器添加JobID和PSN两个字段记录当前占用的分组，并检查哈希冲突。当哈希冲突发生时，后抵达聚合器的分组被继续发送到服务器端处理，完成聚合操作。如果是Reduce操作，该服务器是接收者；如果是AllReduce操作，该服务器是一个额外的服务器。整个冲突处理的过程叫做回退（fallback）。ATP[4]、A2TP[5]采用的就是这种方案。

两种模式对比

        针对多作业场景，独享模式适合给每个作业分配不受其他作业影响的交换机内存空间，并保障作业在特定的时间内完成。但是，独享模式会导致资源利用率较低，一个作业需要2倍于窗口的内存空间，且作业计算时不发送数据，交换机资源闲置。共享模式的资源利用率更高，多个作业以每个往返时间（Round-Trip Time，RTT）的粒度快速地占用和释放聚合器，且一个作业加入时可以动态地占用交换机聚合器；结束时，其他作业可以快速填补该作业不再使用的聚合器。但是，作业之间相互抢占聚合器资源，相互干扰，无法提供可保证的性能。

应用场景与研究

        在生产场景中，通常部署一套集群作为基础设施，多用户提交多作业，运行时共享集群。提升集群的利用率有重要意义，可以提升所有用户作业的整体效率，让集群承载更多的用户作业。在具体场景中，用户对作业的需求不同，例如尽力而为或者有截止时间；集群资源管理的自由度或者权限也不同，例如路由、作业放置、资源调度等。针对具体场景的不同需求，我们可以灵活地组织并改进交换机内存分配模式，从而提升不同的共享集群多作业场景下的整体效率。

面向尽力而为的多作业场景的资源抢占式调度策略

        在共享模式下，在网聚合系统尽力而为地加速所有作业；系统公平对待所有的作业，以先到先服务的方式允许所有作业占用聚合器直至聚合完成。但是，在生产环境中，并非所有作业的所有流量对系统整体效率都有相同的贡献。比如，在一个作业中，神经网络的前层张量的聚合完成后可以快速开启下一轮的前向计算，因此优先传送前层张量比传送后层张量收益大；对于不同神经网络类型的作业，通信占比大的神经网络从在网聚合中获得的收益更大；所有的作业启动和结束时间不同，让剩余运行时间少的作业获得更多资源从而提前完成，可以有效地降低平均作业完成时间。针对这些观察，我们在在网聚合系统中引入了抢占机制[6]。

       每个分组中携带一个优先级字段，每个聚合器中也添加优先级字段。当分组更新聚合器时，一并填入自己的优先级。当聚合器被一个张量的分组占据，且另一个张量的分组抵达时，分组和聚合器进行优先级比较，如果分组优先级高，则分组占据聚合器，并将聚合器内容踢出至参数服务器，否则，分组透传至参数服务器进行处理。

       每个作业在运行过程中，动态地计算自己张量的优先级，并编码到张量分组的头部字段中。优先级的计算考虑该张量在神经网络中的位置、该作业的通信占比和该作业的剩余完成时间这三个因素的乘积。

        通过在交换机中支持基于优先级的资源调度，以及在端侧赋予张量分组优先级，对系统效率贡献大的分组有更大的机会占据聚合器，从而提升尽力而为的多作业场景下的系统整体效率。

面向有截止时间的混合作业场景的资源调度策略

       在部分生产场景中，作业提交者对作业完成时间有要求，例如新产品上线、限时竞赛等。这种情况下，集群资源的调度应该首先满足有截止时间要求的作业，同时，终止不能满足的作业并通知使用者。针对这种场景，我们设计了交换机资源向多作业分配的机制[1]。

        系统采用独享内存模式下的内存分配，为每一个作业指定自己的内存空间。但是，系统允许多个作业各自的内存空间有重叠，因为不同作业的内存可能有重叠部分，系统采用随机寻址加回退机制的方案，每个分组的聚合器寻址方法为：aggregator.index=Hash(JobID, PSN)%Size+Offset。

        运行过程中，集群控制器周期性地调度每个作业的内存分配。首先，将所有的作业按照截止时间的紧急程度排序；其次，从截止时间最早的作业开始，估算其需要的资源；最后，按照估计的资源需求将交换机内存分配给每个作业，分配结束后，剩余的资源被剩余的作业（无截止时间要求和暂时不能满足的作业）共享。

        资源估算方法根据截止时间和收敛剩余的轮次估计每轮的时间，每轮时间扣除测量的计算时间得到每轮的通信时间；根据每轮通信时间和模型大小计算每轮通信需要的带宽；根据估算的带宽和测量得到的RTT计算时延带宽积，即滑动窗口大小；交换机内存需求量就是窗口大小。

       通过将交换机内存以独享的模式分配给有截止时间要求的作业，同时将剩余的内存共享给尽力而为的作业，可以有效地提升截止时间的满足率和整体的多作业执行效率。

面向尽力而为的多作业场景的作业放置策略

        上述策略针对的是作业放置完毕后，如何在运行时分配交换机内存的问题。当集群规模很大时，在集群上部署多个作业会面临如何有效放置作业的问题。一个作业需要消耗集群中的GPU、带宽、交换机内存等多维度资源，不合理的放置会导致网络局部的某种资源紧张而其他资源闲置，集群资源利用率无法达到最优。设计一个有效的作业放置策略涉及两个方面的问题：如何评估一个放置方案的价值，以及如何在所有候选放置方案中高效地搜索最优或较优解[7]。

       假设一个作业放置到已经有若干个作业的集群中，那么已有的作业会竞争网络资源并收敛到一个稳态。估计这个稳态即可估计出该作业能够获取的资源，从而评估该放置方案的有效性。我们改进了注水算法（water-filling），用以估计这个稳态。注水算法先将所有作业和资源标记为活跃状态，然后从零开始逐步向所有作业注入相等的流量，每个作业也将消耗其流经的资源。消耗的带宽等于注入的流量，而交换机消耗的内存等于流量与RTT的乘积。部分资源承载的作业多，因而其叠加的消耗量也多，流量注入过程中，某个资源将先到达容量上限。注水算法将该资源和资源上的作业冻结，然后继续向剩余的活跃作业注入流量，直至所有作业均冻结。注水算法将输出稳态时每个作业的资源使用情况。我们根据集群稳态时每个服务器剩余的网络带宽和GPU数量对服务器进行价值评估。

       进一步，系统需要搜索整个集群中高价值服务的组合来选择当前作业的最佳位置。候选方案的数量是阶乘级的，穷举搜索在时间上不可行。考虑到多服务器的放置方案是由若干个服务器数量少的子方案组合而成的，因而可以用动态规划的方法，从小到大地寻找高价值放置方案，将搜索时间降为多项式级别。

       通过精确地评估每个放置方案的价值和高效地搜索放置方案的候选空间，我们可以有效地解决在大规模在网聚合集群中放置多作业的问题，使全局资源得到均衡、高效的利用。

未来工作展望

       在网聚合技术从底层改变了网络的体系结构，目前能够做到功能上与以太网兼容。但是，网络原本分层的体系结构中，每层功能的设计有其原本的目标。加入在网聚合技术后，这些功能都需要重新考虑，其中存在很大的研究空间。这里列举两个值得在在网聚合背景下深入探索的网络经典问题。

        1.计算机网络原本的路由系统主要面向单播设计，而在网聚合的作业有多个端侧，且构成树状结构。如何在一个拓扑中为多棵树做路由系统，是一个新的挑战[8]。

        2.计算机网络的拥塞控制以流为粒度，通过线性增乘性减的方法，最后收敛到带宽饱和与每流公平的稳态。多个在网聚合流共存时，如何调整拥塞控制，让在网聚合的树状流公平且饱和，也是一个新的挑战[5]。

结束语

       在网聚合是一种能够有效加速分布式机器学习的技术，学术界和工业界都提出了若干原型系统。在网聚合消耗交换机内存资源，在多作业共享集群的场景下，有效地将交换机内存分配给多个作业，可以显著提升多作业整体执行效率和系统资源利用率。结合在网聚合技术重新思考网络管理中的问题，也是一个很大的研究空间。  ■

参考文献

[1] Zhao B, Liu C, Dong J, et al. Enabling Switch Memory Management for Distributed Training with In-Network Aggregation[C]// 2023 IEEE Conference on Computer Communications (INFOCOM), 2023.

[2] Sapio A, Canini M, Ho C Y, et al. Scaling Distributed Machine Learning with In-Network Aggregation[C]// 18th USENIX Symposium on Networked Systems Design and Implementation (NSDI 21), 2021.

[3] Liu S, Wang Q, Zhang J. et al. In-Network Aggregation with Transport Transparency for Distributed Trainings[C]// Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems. 2023.

[4] Lao C, Le Y, Mahajan K, et al. ATP: In-network Aggregation for Multi-tenant Learning[C]// 18th USENIX Symposium on Networked Systems Design and Implementation (NSDI 21). 2021:. 741-761.

[5] Li Z, Huang J, Li Y, et al. A2TP: Aggregator-Aware In-Network Aggregation for Multi-Tenant Learning[C]// Proceedings of the Eighteenth European Conference on Computer Systems. 2023: 639-653.

[6] Wang H, Qin Y, Lao C, et al. Preemptive Switch Memory Usage to Accelerate Training Jobs with Shared In-Network Aggregation[C]// Proceedings of The 31st IEEE International Conference on Network Protocols. 2023.

[7] Zhao B, Xu W, Liu S, et al. Training Job Placement in Clusters with Statistical In-Network Aggregation[C]// Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems. 2024.

[8] 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, 31(5): 2267-2280.

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