本帖最后由 中国计算机学会 于 2024-4-19 14:59 编辑
摘要—本文旨在通过系统性地分析分布式AI模型训练面临的挑战与在网计算的收益原理，向读者全面介绍在网计算这一技术，并回顾在网计算相关技术在业界的研究进展，借此对在网计算未来的发展方向提出展望。 刘 硕（华为公司）
徐 捷（华为公司）
关键词 ：在网计算 大模型 分布式训练
AI大模型分布式训练的流量分析 分布式训练 AI模型训练的算力需求大致分为3个阶段[1]。AI早期（~2012年），由于参数量规模不大，摩尔定律尚可满足需求。2012年，AlexNet开启了深度学习时代，根据OpenAI的分析，2012~2018年，算力每3.4个月翻一倍。2018~2022年，算力每5.7个月翻一倍，已经远超摩尔定律定义的能力。以ChatGPT为首的生成类模型取得技术突破后，大模型时代正式到来，AI训练的算力需求远超单芯片能力，分布式训练是必然结果。 如今，大模型参数量已迈入万亿规模。2021年，百度发布悟道大模型，参数量达到了17500亿，超越了Google Switch Transformer的16000亿，远超OpenAI GPT-3的1750亿。尚未披露细节的GPT-4，据推测参数量也达到了万亿规模。阿里参与打造的八卦炉项目已经瞄准了百万亿的参数规模，如图1所示。大模型参数量的不断增加，意味着通信开销逐渐成为整个集群的性能瓶颈。
集合通信流量分析 分布式训练必然引入网络通信开销。当前大模型分布式训练的并行策略主要分为张量并行、数据并行、流水线并行和混合专家（Mixture of Experts，MoE）并行[2]，如图2所示。不同的切分策略对应不同的通信操作。这里结合切分策略对通信流量的类型进行简要介绍。
1.张量并行产生All-Reduce集合通信流量：张量并行由模型参数量驱动，随着模型尺寸的不断增大，需要把一个完整模型分布在不同的计算卡上。一个张量（tensor）沿特定维度被切分成若干部分，分布在不同的计算卡上分别计算，不同计算卡共用同一份输入，最终将各个计算卡的计算结果通过All-Reduce操作在不同卡间进行汇聚。张量并行产生的All-Reduce流量在前向计算和后向传播过程中都会发生。 2.数据并行产生All-Reduce集合通信流量：数据并行由训练数据集驱动，随着训练数据的不断增加，需要把不同训练样本分布在不同的计算卡上。每个计算卡基于不同的输入数据集训练会产生不同的梯度结果。为保证模型的最终收敛，需要将不同计算卡的梯度通过All-Reduce操作进行汇聚。数据并行产生的All-Reduce流量主要发生在后向传播过程中。 3.流水线并行产生Send/Recv流量：流水线并行的核心思想是把模型按层切割成若干部分，每张计算卡负责一个部分。在前向传递过程中，每张卡基于Send操作将中间的激活结果传递给下一个阶段；在后向传递过程中，每个计算卡将输入张量的梯度传回给前一个流水线阶段，也即每个阶段基于Recv操作获取张量的梯度。 4.混合专家并行产生All-to-All流量：混合专家并行由降低算力需求驱动，只有小部分网络用于计算任何一个输入的输出，降低对计算量需求的同时，扩大了模型参数量，具有更高的稀疏性。但是由于每个专家只负责部分特征的计算，需要专家之间基于All-to-All操作传递各自的特征计算结果。 以GPT-3为例，按16路数据并行，8路张量并行，8路流水线并行，批量大小（batch size）取1536，统计各个并行策略调用通信原语的数据传输量，如表1所示，All-Reduce原语的数据传输量占比是最高的。
All-Reduce集合通信的算法 All-Reduce算法有各种各样的实现，典型的有Scatter-Reduce + All-Gather和Reduce + Broadcast两种模式。下文对以Scatter-Reduce + All-Gather模式实现的Ring All-Reduce算法[3]进行简要介绍。 Ring All-Reduce算法
Ring All-Reduce算法以环（Ring）的形式完成Scatter-Reduce和All-Gather阶段的数据传递，是带宽最优的算法。 假设有n张计算卡参与计算任务，传输数据量为M。把M分成n份数据块，每一次All-Reduce操作发送的数据量为M/n。在Scatter-Reduce阶段，每次迭代，计算卡都向其下游节点发送其中一份数据块，同时接收其上游节点发来的数据块。这样可以最大程度地利用计算卡的双向带宽。计算卡将收到的数据块与本地对应的数据块求和。不同迭代轮次，发送的数据块不同。经过n－1轮迭代，每张卡上都会有一个特定的数据块为全局的求和结果，如图3所示。
在All-Gather阶段，计算卡不再对接收到的数据块求和，而是替代自己本地对应的数据。又经过n－1轮迭代，每张卡上的所有数据块均为全局的求和结果，如图4所示。综上所述，一次All-Reduce的通信开销为：
Ring All-Reduce时延 = 2(n－1)α + 2(n－1)/n×M/B （1） 其中α为一次发送的开销，包括端侧协议栈处理时延；B为传输带宽。 由公式（1）可知，Ring All-Reduce算法下，节点间需要频繁交互，共2(n－1)轮次，因此当n越来越大，即参与计算的节点数越来越多时，发送开销越来越大。此外，在2(n－1)轮次的数据收发中，单张计算卡传输的数据量接近2M。 在网计算
在网计算利用网络交换机的中心拓扑位置优势，随路对All-Reduce流量进行聚合，如图5所示，以避免过多的收发次数。在网计算技术通过降低节点交互频率，一方面降低处理时延开销，另一方面减少大约50%的数据传输量，从而加速All-Reduce操作。表2对比了在网计算与不同All-Reduce算法的通信开销。首先，发送次数降低到1次时，相比Ring算法，在网计算的静态时延至少减少50%（n=2）。其次，因为数据经过交换机之后已经完成随路聚合，所以单张计算卡只需要发送所有的数据量到交换机，得到的就是全局的求和结果。数据传输量从原来的2M下降到M，传输时延减少50%。
相比其他的HD（Halving Doubling）算法[4]，在网计算同样将节点交互次数从2log2n降低为1，处理时延也至少降低50%。数据传输量也从原来2M下降到M，传输时延减少50%。
在网计算技术的研究进展
工业界部署情况 在网计算在工业界的部署，最早出现在超算（High Performance Computing，HPC）集群，用于卸载集合通信操作，释放CPU算力，如IBM的BlueGene-Q[5]、富岳超算的Tofu[6]、英伟达（NVIDIA）的SHARP v1[7]等，如表3所示。NVIDIA的SHARP是目前唯一大规模商用的在网计算加速AI训练解决方案。下面着重介绍SHARP，并将其与华为的在网计算解决方案进行对比。
如表4所示，SHARP v1最早用于超算场景的All-Reduce加速。虽然超算场景中一次All-Reduce的数据量较小，但频次很高，因此All-Reduce时延不可忽略。针对该场景，NVIDIA在无限带宽（InfiniBand，IB）网络交换机转发逻辑旁挂载目标通道适配器（Target Channel Adapter，TCA）模块，专门用于接收All-Reduce流量。实际系统运行过程中，交换机显式地参与计算任务，对端侧计算卡可见。计算卡直接将All-Reduce数据发送至IB交换机TCA端口，接收后的流量交由专门处理聚合运算的算术逻辑单元（Arithmetic Logic Unit，ALU）进行汇聚后，将最终结果广播给计算卡。由于超算All-Reduce数据量不大，采用这样的方式对现行交换架构影响较小，带宽压力不大，但无法应对AI场景下的大消息流。随后，SHARP相继推出了v2[8]和v3版本，采用在交换机端口处汇聚的方式应对AI流量对带宽造成的冲击。这样做的代价是交换机端口需要配备TCA模块，处理IB网络协议栈，同时还需要配备一定的缓存和计算能力。
英伟达收购迈络思（Mellanox）后，为了增强自家GPU的带宽能力，将原先用于IB网络的SHARP在网计算技术移植到NVLink网络中。例如，当前最新发布的NVIDIA DGX H100服务器[9]在NVSwitch中部署了SHARP在网计算能力。相比上一代A100 GPU“2N发送，2N接收”的All-Reduce开销，H100+NVLink SHARP利用交换机在网计算，实现了“N+1发送，N+1接收”，有效带宽提升1倍。如表5所示，上一代DGX A100没有部署在网计算能力，其All-Reduce带宽为150 GB/s，只能利用A100 GPU单向满带宽300 GB/s的50%；SuperPod场景（DGX A100*32，共256张A100卡），All-Reduce带宽为100 GB/s，只有GPU带宽的33%。相比A100，在SHARP的加持下，DGX H100无论在服务器内还是跨服务器的SuperPod集群，All-Reduce带宽都可以用满H100 GPU单向带宽450 GB/s，分别有3倍和4.5倍的带宽提升。 学术界研究进展 除工业界研究火热以外，在网计算近年来在学术界也引发了研究热潮，在网计算加速AI训练的学术文章在计算机系统顶会上屡见不鲜。学术界关于在网计算加速AI训练，主要以基于可重构匹配表（Reconfigurable Match Table，RMT）架构[12]的Intel Tofino[13]可编程网络交换芯片为平台，聚焦在如何基于受限资源不影响线速转发的前提下，实现在网计算功能。主要有以下两大流派： 1.规定任务上限，对交换机资源做预留，控制面逻辑简单，但资源利用率不高，例如SwitchML[14]（NSDI’21）、SHARP[8]（ISC’20）等。 2.不规定任务上限，基于哈希（hash）进行动态资源复用，哈希冲突后回退传统端侧All-Reduce模式。资源利用率高，但控制面处理逻辑复杂，例如ATP[15]（NSDI’21）等。 除此之外，一些研究工作也利用其他平台，如FPGA、NP芯片等，进行在网计算相关的研究。如iSwitch[16]（ISCA’19）利用FPGA加速强化学习训练，TRIO[13]（SIGCOMM’22）利用Jupiter的Trio NP芯片进行在网计算的实践等。同时，也有研究者在关注引入在网计算对现有网络架构的影响，如PANAMA[18]（MLSys’21）讨论了如何解决拥塞控制问题，TRIO[17]（SIGCOMM’22）讨论了Straggler问题，等等。 华为关于在网计算的实践 笔者所在的华为2012网络技术实验室在网计算团队提出了一种传输层透明的网内计算原语NetReduce[19]，用来加速AI模型训练。与SHARP基于IB网络采用TCA截断端到端连接的方式不同，NetReduce无须在网络中部署额外的端口网卡，而是基于以太网，有更好的可扩展性，如图6所示。交换机不终结计算节点的端到端连接，All-Reduce流量经过时，采用中间人（man-in-the-middle）的方式收集相关数据并聚合，替换原来报文中的有效载荷（payload），报文状态还是交由端侧维护，无须处理网络协议栈，做到了传输层透明。
具体来说，在主机侧使用RoCE（RDMA over Converged Ethernet）v2作为传输层协议，传递梯度消息并接收聚合结果。在网络侧，利用交换机执行聚合操作，但保留数据传输连接。这样可以从网内聚合和硬件支持RoCE两方面获得性能收益：一方面网内聚合带来传输数据量减少、释放带宽、线性扩展等优势；另一方面，通过硬件卸载RoCE协议处理能够降低CPU开销，带来高吞吐、低延迟的网络IO性能。
NetReduce技术针对通信密集型模型的数据并行场景有明显的优化效果。如图7所示，在数据并行All-Reduce通信开销不被计算掩盖的前提下，在网计算对通信密集型模型AlexNet和VGG-16分别有68.8%和50.7%的端到端性能提升。而对于计算密集型模型ResNet-50，收益会有所下降，端到端性能提升为15.1%。 但是针对大模型，数据并行All-Reduce通信过程通常会与计算互相重叠，通信开销部分隐藏在计算开销内。因此放到端到端的范围来看，在网计算减少数据并行All-Reduce通信开销带来的收益有可能被计算掩盖。而张量并行All-Reduce在大模型训练当中同样不可或缺，而且由于计算与通信之间通常没有重叠，在网计算降低通信开销的收益往往能够全部体现在端到端中。因此在网计算对于大模型的收益主要来自对张量并行All-Reduce的优化。基于华为昇腾计算卡的仿真结果显示，在网计算对千亿参数量模型带来的端到端收益在12%左右。
总结 2023年是以GPT-3.5为首的生成式AI大模型爆发的元年，随着未来大模型参数量不断增加，网络通信在模型训练中的瓶颈将逐渐显现。在网计算将计算节点间的交互复杂度从O(n)/O(log n)降低为O(1)，数据传输量减少大约50%，加速All-Reduce集合通信操作，从而端到端提升AI模型训练吞吐，在业界引起了广泛的关注讨论与研究热情。SHARP等相关技术的实践与落地，也意味着在网计算技术能够在客户层面带来巨大的商业价值。但在网计算领域还有许多尚未解决的问题与挑战，例如，如何降低落地芯片的代价与开销，如何在异构网络拓扑下实现高效的在网计算，以及是否还能够加速除All-Reduce以外的其他集合通信原语，等等。相信在工业界和学术界的共同努力下，在不久的将来，在网计算能够成为网络领域的基础技术，得到广泛应用。 ■ 参考文献： [1] Sevilla J, Heim L, Ho A, et al. Compute trends across three eras of machine learning[C]// 2022 International Joint Conference on Neural Networks (IJCNN), 2022:1-8. [2] Narayanan D, Shoeybi M, Casper J, et al. Efficient large-scale language model training on GPU clusters using megatron-lm[C]// Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, 2021:1-15. [3] Gibiansky A. Bringing HPC techniques to deep learning. Baidu Research, Tech. Rep, 2017. [4] Rabenseifner R. Optimization of collective reduction operations[C]// Computational Science-ICCS 2004: 4th International Conference, 2004:1-9. [5] Chen D, Eisley N, Heidelberger P, et al. The IBM Blue Gene/Q Interconnection Fabric[J]. IEEE Micro, 2012,32(1):32-43. [6] Ajima Y, Inoue T, Hiramoto S, et al. Tofu: Interconnect for the K computer[J]. Fujitsu Sci. Tech. J, 2012, 48(3): 280-285. [7] Graham R L, Bureddy D, Lui P, et al. Scalable hierarchical aggregation protocol (SHArP): A hardware architecture for efficient data reduction[C]// 2016 First International Workshop on Communication Optimizations in HPC (COMHPC), 2016:1-10. [8] Graham R L, Levi L, Burredy D, et al. Scalable hierarchical aggregation and reduction protocol (SHARP) TM streaming-aggregation hardware design and evaluation[C]// High Performance Computing: 35th International Conference, 2020:41-59. [9] Ishii A, Wells R. The Nvlink-Network Switch: Nvidia’s Switch Chip for High Communication-Bandwidth Superpods[C]// 2022 IEEE Hot Chips 34 Symposium (HCS), 2022:1-23. [10] Lakhotia K, Petrini F, Kannan R, et al. In-network reductions on multi-dimensional HyperX[C]//2021 IEEE Symposium on High-Performance Interconnects (HOTI). IEEE, 2021: 1-8. [11] Yang M, Baban A, Kugel V, et al. Using trio: juniper networks' programmable chipset-for emerging in-network applications[C]//Proceedings of the ACM SIGCOMM 2022 Conference. 2022: 633-648. [12] Bosshart P, Gibb G, Kim H S, et al. Forwarding metamorphosis: Fast programmable match-action processing in hardware for SDN[J]. ACM SIGCOMM Computer Communication Review. 2013, 43(4), 99-110. [14] Sapio A, Canini M, Ho C Y, et al…. Scaling distributed machine learning with in-network aggregation[OL]. arXiv preprint arXiv:1903.06701, 2019. [15] Lao C, Le Y, Mahajan K, et al. ATP: In-network Aggregation for Multi-tenant Learning[C]// NSDI’21. 2021:741-761. [16] 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. 2019,6: 279-291. [17] Yang M, Baban A, Kugel V, et al. Using trio: juniper networks' programmable chipset-for emerging in-network applications[C]// Proceedings of the ACM SIGCOMM 2022 Conference. 2022, 8: 633-648. [18] Gebara N, Ghobadi M, Costa P. In-network aggregation for shared machine learning clusters. [C]//Proceedings of Machine Learning and Systems 3 (MLSys 2021). 2021: 829-844. [19] Liu S, Wang Q, Zhang J, et al. In-Network Aggregation with Transport Transparency for Distributed Training[C]// Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems. 2023,3: 376-391.
版权声明：中国计算机学会(CCF)拥有《中国计算机学会通讯》(CCCF)所刊登内容的所有版权，未经CCF允许，不得转载本刊文字及照片，否则被视为侵权。对于侵权行为，CCF将追究其法律责任。
| 
京公网安备 11010802032778号 )
发表于 2024-4-12 16:57:21