UB-Mesh: a Hierarchically Localized nD-FullMesh Datacenter Network Architecture
作者/机构: Heng Liao, Bingyang Liu, Xianping Chen, Zhigang Guo, Chuanning Cheng, Jianbing Wang, Xiangyu Chen, Peng Dong, Rui Meng, Wenjie Liu, Zhe Zhou*, Ziyang Zhang, Yuhang Gai, Cunle Qian, Yi Xiong, Zhongwu Cheng, Jing Xia, Yuli Ma, Xi Chen, Wenhua Du, Shizhong Xiao, Chungang Li, Yong Qin,Liudong Xiong, Zhou Yu, Lv Chen, Lei Chen, Buyun Wang, Pei Wu, Junen Gao,Xiaochu Li, Jian He, Shizhuan Yan, Bill McCOLL (Huawei)
A1 主要贡献
随着大型语言模型(LLM)的持续扩展,其对计算能力和带宽的需求也随之升级。本文针对下一代AI数据中心面临的大规模(R1)、高带宽(R2)、成本效益(R3)和高可用性(R4)四大挑战,提出了一个名为UB-Mesh的新型AI数据中心网络架构。
该研究的核心贡献是:
1. 提出了一种分层局部化的nD-FullMesh网络拓扑:与传统数据中心提供对称节点间带宽的Clos架构不同,UB-Mesh利用LLM训练的数据局部性,优先采用短距离直接互连,以最小化数据移动距离和减少交换机使用。这种设计从1D(板内)到nD(跨机柜、跨楼层)递归地构建全连接拓扑,能够分层分配带宽,满足LLM训练中不同并行策略(如张量并行、数据并行)对通信带宽的不同需求。
2. 设计并实现了一套完整的硬件构建模块和统一总线(UB)技术:为了将理论拓扑付诸实践,本文设计了包括NPU、CPU、低基数交换机(LRS)、高基数交换机(HRS)和NIC等在内的一系列硬件组件。这些组件通过一种新颖的统一总线(UB)技术互连,实现了灵活的IO带宽分配和硬件资源池化,并简化了跨层优化。
3. 开发了针对UB-Mesh架构的网络系统优化机制:
* 全路径路由(APR):为有效管理数据流量并充分利用直接连接链路的带宽,提出了APR机制。它结合了源路由、结构化寻址和无死锁流控技术,实现了自适应路由并最小化了转发开销。
* 拓扑感知的性能增强:提出了拓扑感知的集合通信算法(如多环AllReduce和多路径All-to-All)和并行化策略,以优化带宽利用率。
4. 实现了高可用性和可靠性设计:
* 64+1备份设计:每个机柜包含一个额外的备用NPU,当系统内有NPU发生故障时,可激活备用NPU以保证LLM训练任务不中断。
* 快速故障恢复:路由系统通过新颖的直接通知技术,在链路故障时能快速恢复。
综合评估表明,与传统的无超分Clos网络相比,UB-Mesh架构将高基数交换机使用量减少了98%,光模块使用量减少了93%,系统级成本效益提高了2.04倍,网络可用性提高了7.2%,并在各种LLM训练任务中实现了95%以上的线性度。
(c) Clos与UB-Mesh的特性比较
图1: 传统Clos数据中心架构与UB-Mesh的比较
A3 背景知识与动机
2.1 LLM训练中的“通信墙”
LLM训练是史上规模最大、计算和通信最密集的并行计算应用。遵循所谓的扩展定律(scaling laws)【【3,Fire-Flyer AI-HPC: A Cost-Effective Software-Hardware Co-Design for Deep Learning,2024】,【5,Unified scaling laws for routed language models,2022,International conference on machine learning】,【13,Scaling laws for neural language models,2020,arXiv】】,LLM通过扩展模型参数和训练数据来提升性能,这导致需要越来越多的AI加速器(NPU/GPU)才能在合理时间内完成训练。例如,开源的LLAMA-3.1模型在16K个GPU上进行训练【【8,The llama 3 herd of models,2024,arXiv】】,而下一代LLM模型正在使用10万个GPU进行训练【【24,100,000 H100 Clusters: Power, Network Topology, Ethernet vs InfiniBand, Reliability, Failures, Checkpointing,2024,Semianalysis】,【29,Elon Musk’s xAI to Use 100,000 NVIDIA H100 GPUs to Train Grok-3 AI Model, Grok-2 Launch in August,2024,WCCFtech】】。标准的训练过程包括前向传播、后向传播和优化器更新。为了充分利用分布式计算能力,LLM训练通过各种并行策略将数据、模型和激活值分割到数万个NPU上。在每次迭代中,NPU之间频繁交换数据。随着训练系统规模的扩大,数据移动成为系统中最昂贵的部分【【9,RDMA over Ethernet for Distributed Training at Meta Scale,2024,ACM SIGCOMM 2024 Conference】,【19,Alibaba HPN: A Data Center Network for Large Language Model Training,2024,Traffic】】。如果没有强大的NPU间通信能力,训练过程很容易受到“通信墙”的瓶颈限制。
2.2 LLM训练中数据流量的局部性
LLM训练通常涉及多种并行技术,如图2所示。
* 张量并行(Tensor Parallelism, TP):TP以行或列的方式切分模型层,并将子层放置在多个并行计算的NPU上【【25,Megatron-lm: Training multi-billion parameter language models using model parallelism,2019,arXiv】】。它主要涉及AllReduce操作来合并分布式的部分结果。
* 序列并行(Sequence Parallelism, SP):SP(在某些论文中也称为上下文并行)通常用于将序列分割到多个NPU上以实现并行处理。SP依赖于RingAttention【【15,Ring Attention with Blockwise Transformers for Near-Infinite Context,2023,arXiv】】等技术,并可使用AllGather操作来收集不同NPU上的部分结果。
* 专家并行(Expert Parallelism, EP):对于采用专家混合(MoE)技术的LLM模型,密集的MLP层被MoE层取代,每层包含多个“专家”。专家在执行过程中被稀疏激活。EP将专家分布在不同的NPU上。输入令牌通过All2All通信动态发送到目标专家。
* 流水线并行(Pipeline Parallelism, PP):与TP切分每一层不同,PP将不同的层分布到多个设备上,并以流水线方式执行前向和后向传播。PP涉及用于跨层传输激活值的低开销P2P通信,但需要高效的调度算法【【11,Gpipe: Efficient training of giant neural networks using pipeline parallelism,2019,Advances in neural information processing systems】,【16,PipeDream: Generalized pipeline parallelism for DNN training,2019,ACM symposium on operating systems principles】】来最小化气泡。
* 数据并行(Data Parallelism, DP):DP在多个NPU上复制模型和优化器状态。每个副本并行处理一部分输入批次。在训练过程中,需要AllReduce操作来同步梯度。
图2: LLM训练中的并行性
表1: LLM训练中的数据流量分析
关键观察:通信强度是分层的并表现出强局部性。根据我们对内部MoE-2T模型的分析(见表1),并非所有并行技术产生的数据流量都相同。具体来说,TP和SP约占总流量的97%,而其余并行技术通常产生不到2%的总流量。其他模型架构可能数据流量分布略有不同,但也表现出强烈的局部性【【28,Rail-only: A Low-Cost High-Performance Network for Training LLMs with Trillion Parameters,n.d.】】。因此,架构设计应优先考虑采用分层的方法来提供网络带宽。
2.3 数据中心网络架构
设计大规模LLM训练集群的关键考虑之一是如何将众多NPU资源组织成一个统一的分布式系统。传统数据中心和超算系统探索了多种拓扑,但它们可能不适合大规模LLM训练。图3展示了一些流行的拓扑结构:
* Clos:传统数据中心通常采用Clos架构,使用两层或三层交换机对称连接NPU/CPU,为各种流量模式提供高性能和适应性。然而,其主要缺点是成本高昂,因为它大量使用高性能交换机和光模块。
* 3D Torus:与依赖交换机互连的Clos不同,3D Torus直接连接相邻的NPU。虽然降低了成本,但3D Torus提供的NPU间带宽较低,且不太适应复杂的通信模式,如在MoE模型中普遍存在的All-to-All通信【【20,DeepSpeedMoE: Advancing Mixture-of-Experts Inference and Training to Power NextGeneration AI Scale,2022,Proceedings of the 39th International Conference on Machine Learning】】。
* Dragon-Fly:Dragon-Fly拓扑【【14,TechnologyDriven, Highly-Scalable Dragonfly Topology,2008,Proceedings of the 35th Annual International Symposium on Computer Architecture (ISCA ’08)】】将交换机组织成互连的组,组内交换机全互连。这些组再通过长距离线缆连接。该架构在传统数据中心和HPC环境中降低了维度。DF比Clos便宜,但由于需要完整的NPU-交换机带宽,成本仍然很高,并且其架构在P2P/AllReduce等LLM训练流量场景下表现不佳。
* Fugaku Tofu:Fugaku Tofu【【2,Tofu: A 6D mesh/torus interconnect for exascale computers,2009,Computer】】为HPC应用提出了一种独特的6D Torus拓扑。然而,与3D Torus类似,它提供的NPU间带宽较低,且不适应复杂的通信模式。
图3: 传统数据中心网络架构
总结而言,如何为大规模LLM训练设计一个高性能且成本效益高的数据中心网络架构仍然是一个悬而未决的问题。
A2 方法细节
3 UB-Mesh 架构
为了满足第1节中提出的R1-R4要求并遵循P1-P3原则,我们提出了UB-Mesh架构。UB-Mesh采用了一种新颖的nD-FullMesh拓扑,最大限度地利用直接电缆互连,从而减少了对昂贵的高带宽交换机和光模块的需求。我们接着介绍了一系列专门为构建UB-Mesh设计的硬件模块。最后,我们将讨论UB-Mesh-Pod(UB-Mesh的具体4D-FullMesh实现)的详细架构设计,其中充分考虑了多种工程约束和权衡。
3.1 nD-FullMesh 拓扑
UB-Mesh引入了一种新颖的nD-FullMesh拓扑,强调在所有网络层级上进行直接互连。如图4所示,该拓扑从相邻节点形成一个1D全连接网格(full-mesh)开始,其中每个节点都与同一层级的其他所有节点相连。扩展此概念,相邻的1D全连接网格内的节点再进行直接互连,创建一个2D全连接网格。这个过程递归地继续,相邻的2D全连接网格形成一个3D全连接网格,以此类推,最终根据需要扩展到一个nD全连接网格。这种灵活的虚拟拓扑可以无缝地映射到各种物理NPU组织形式:1D全连接网格可以部署在单个板卡上,2D全连接网格部署在机柜内,3D全连接网格跨越机柜行,4D全连接网格在一个楼层的机柜组内,甚至5D全连接网格可以跨越整个建筑物,乃至更高维度。
图4: UB-Mesh的nD-FullMesh拓扑及其可能的物理映射
这种分层局部化的拓扑在构建下一代LLM训练数据中心方面具有几个优势。首先,由于nD-FullMesh在每个网络维度上形成了紧密耦合的直接连接域,并能提供逐级超分的带宽,它充分利用了LLM训练中固有的数据局部性和从密集到稀疏的流量模式(见2.2节)。此外,我们可以灵活调整每个维度的每节点带宽分配,以满足未来LLM训练任务的特定需求。例如,如图5-(a)所示,假设存在一个6D-FullMesh拓扑,我们可以通过在NPU的IO模块中分配不同数量的互连链路来调整$D_x$和$D_y$维度的带宽,如图5-(b)所示。
图5: 高维拓扑与灵活带宽分配
表2: 不同类型链路的使用估算
其次,与需要大量交换设备的Clos架构不同,这种直接连接的拓扑为减少传输距离提供了巨大潜力。结合数据放置和集合通信优化(见5.2节),大多数传输(通常涉及TP或SP,如表1所示)可以在0-2跳内完成,从而大大减少了数据移动的开销。
第三,与Clos等其他流行拓扑相比,nD-FullMesh拓扑减少了对高带宽交换机和长距离光互连的依赖。相反,它最大化地利用了短距离直接互连。例如,根据我们的估算,如表2所示,短距离无源电缆占所消耗电缆总量的86.7%。这种设计大幅降低了交换机和光模块的成本,同时也提高了系统可靠性,因为电缆和连接器比光模块更稳定。
最后,与不适合如all-to-all等复杂集合通信操作的3D-Torus拓扑不同,full-mesh拓扑可以高效地支持这些操作,这将在第5节中讨论。
3.2 UB-Mesh的构建模块
所提出的nD-FullMesh拓扑具有若干理论优势,但其实际实现和系统性能的提升带来了新的挑战,需要全面考虑各种工程约束和权衡。为了解决这些问题,我们设计并制造了一系列硬件模块,作为UB-Mesh的基本构建块。这些模块通过一种名为统一总线(Unified Bus, UB)的新型互连机制连接。具体细节如下:
3.2.1 硬件模块。表3列出了UB-Mesh的主要组件。其核心是作为AI计算单元的NPU。NPU配备了两个UB IO控制器,总共提供x72 UB通道。负责执行主机程序的CPU也配备了一个UB IO控制器,提供UB x32 IO。虽然UB-Mesh优先采用直接互连,但交换能力对于满足特定需求仍然至关重要。例如,最好将图5中的$D_z$和$D_w$维度NPU IO进行汇聚,以便于高效的机柜间互连。此外,我们期望CPU-NPU的比率及其绑定关系可以灵活调整,以实现高效的NPU和内存资源池化,这可以通过交换实现。为了在UB-Mesh-Pod中提供成本效益高的交换能力,我们开发了LRS,即轻量级的低基数交换机。它的制造成本低,能有效汇聚机柜间IO带宽并实现CPU-NPU通信。尽管UB-Mesh可以仅使用低基数交换机实现,但实际的网络约束通常将其范围限制在一定规模(UB-Mesh-Pod)。因此,UB-Mesh还包括高基数交换机(HRS),为Pod级别的交换提供UB x512 IO容量。UB-Mesh的详细架构设计和考虑将在3.3节中介绍。
表3: UB-Mesh的主要构建模块
3.2.2 统一总线(UB)互连。连接这些异构硬件模块需要特定的互连技术。如图6-(a)所示,传统的基于GPU的LLM训练系统【【9,RDMA over Ethernet for Distributed Training at Meta Scale,2024,ACM SIGCOMM 2024 Conference】,【19,Alibaba HPN: A Data Center Network for Large Language Model Training,2024,Traffic】】通常采用混合互连:CPU-GPU和CPU-NIC之间使用PCIe,服务器内GPU之间使用NVLink,服务器间使用InfiniBand/RoCE。如图6-(b)所示,与以往的设计不同,我们引入了统一总线(UB)互连。UB同时支持同步操作(如Load/Store/Atomic)和异步操作(如Read/Write/Message)。基于UB,我们只需设计和实现一个单一的UB IO控制器,该控制器可以在CPU、NPU乃至LRS交换机之间复用。UB还提供以下三个主要优势:
* 灵活的IO资源分配:UB互连与特定用例解耦,使得芯片内不同类型的IO资源可以灵活分配,如图5所示。由于它们使用相同的UB链路,NPU间带宽和CPU-NPU带宽也可以根据特定需求灵活调整。
* 硬件资源池化:UB的点对点通信能力使得硬件资源(包括DDR DRAM、CPU、NPU和NIC)的高效池化成为可能。例如,在图6-(b)中,CPU和NPU通过UB互连进行池化,以提高资源利用率。
* 系统优化优势:统一的互连消除了协议转换的需要,显著减少了开销,并简化了驱动程序、通信库和框架等的设计与优化。
图6: 统一总线互连 VS. 传统混合互连
3.3 架构概览
通过组合表3中列出的硬件构建模块并通过统一总线(UB)互连它们,我们展示了UB-Mesh的整体架构设计。如图7-(a)所示,我们遵循提出的nD-FullMesh拓扑实现了一个UB-Mesh-Pod。具体来说,在该架构内,我们在每个机柜内创建了一个2D-FullMesh,并将其扩展到机柜外的另一个2D-FullMesh,最终形成一个4D-FullMesh。我们选择在这一代产品中不将UB-Mesh扩展到5D-FullMesh或更高维度,以在成本效益和灵活性之间取得工程上的平衡,在4D-FullMesh之外仍然采用Clos拓扑。具体细节介绍如下:
图7: UB-Mesh的架构设计
3.3.1 在机柜内实现2D-FullMesh。如图7-(b)所示,每个机柜配置为2D-FullMesh拓扑,集成了多个组件以确保高效的通信和资源利用。机柜的核心由八个NPU板组成,每个板包含八个NPU,如图8-(a)所示。这64个NPU在机柜内互连,形成一个2D全连接网格网络,确保了高的NPU间带宽。值得注意的是,由于UB IO控制器也具备路由能力,每个NPU在此架构中都充当一个路由器,并支持间接路由。除了NPU,机柜还包括专用的CPU板,如图8-(c)所示。与CPU和NPU位于同一板上的传统设置不同,这里它们是分离的。CPU通过交换机连接到NPU,从而实现了灵活的CPU/NPU比率,并支持CPU/NPU/DDR资源池化以提高资源利用率。机柜还配备了多个背板交换机,用于管理机柜内和机柜间的连接。这些交换机采用低基数设计,称为LRS,以降低成本同时确保设备间的无阻塞通信。如图7-(b)所示,18个LRS全连接形成一个交换平面,其中两个LRS用于CPU和备用NPU,八个用于常规NPU,八个用于机柜间连接。这些背板交换机总共输出四个UB x256 IO。
图8: UB-Mesh机柜的硬件实现
3.3.2 64+1高可用性设计。如第1节所强调,确保系统可用性是大规模AI数据中心面临的一个关键挑战。系统中的单个NPU可能会出现意外故障。虽然软件系统可以检测到这些故障并使用剩余的健康NPU重启训练,但由于计算能力和系统带宽的降低,系统性能将受到显著影响。为了提高系统可用性,我们引入了一种特殊的64+1设计:如图7-(b)和图8-(b)所示,系统由64个常规NPU和一个额外的备用NPU组成。这个备用NPU通过LRS连接到64个常规NPU。在任何NPU发生故障时,备用NPU会被策略性地利用。如图9所示,当NPU-3发生故障时,管理系统会激活备用NPU(图中的node-B)来替换NPU-3。与NPU-3相关的原始直接连接链路被重定向。例如,路径5-3被重定向为路径5-LRS-B。尽管这种策略将原来的直接连接变为了一跳路由,略微增加了传输延迟,但它远优于简单地屏蔽NPU-3并在剩余的七个NPU上运行任务。在机柜内带宽是主要考量因素的LLM训练场景中,增加的延迟对整体训练性能的影响可以忽略不计。
图9: 通过启用备用NPU实现容错
3.3.3 在UB-Mesh-Pod中扩展到4D-FullMesh。在一个机柜内,我们将64个NPU组织成一个2D全连接网格拓扑。在机柜之外,我们实现了另一个2D全连接网格来组织机柜,形成一个名为UB-Mesh-Pod的4D-FullMesh架构。如图7-(a)和(c)所示,一行中的四个相邻机柜被组织成一个紧密耦合的1D全连接网格,其中机柜通过每个机柜中的LRS端口直接互连。在每个机柜列中,机柜也直接连接,从而创建一个16机柜的2D全连接网格拓扑。此配置中的每个链路代表一个UB x128 IO,如图8-(d)所示。我们沿两个维度连接四个相邻机柜来构建机柜间全连接网格,因为这是考虑有源电缆可达范围的最佳点。由于每个机柜有64个NPU,每个Pod有16个机柜,一个4D-FullMesh UB-Mesh-Pod总共包含1024个NPU。
3.3.4 UB-Mesh-SuperPod及更高扩展。拥有1K规模的UB-Mesh-Pod后,我们进一步构建UB-Mesh-SuperPod,它可以容纳多个UB-Mesh-Pod。考虑到当前云场景中,中小型LLM训练工作负载可能不会消耗整个SuperPod,我们选择在Pod级别的互连中采用对称的Clos拓扑,而不是继续使用全连接网格。这样的设计允许云管理者根据消费者的需求灵活划分SuperPod,并保证每个划分域内的全带宽。如图7-(c)所示,我们使用高基数Pod交换机(HRS)来连接SuperPod中的每个机柜,规模可扩展至8K个NPU。最后,SuperPod中的机柜也通过UB交换机(方案a)或位于CPU板上的NIC(方案b)连接到大规模的数据中心网络(DCN)。DCN域通常支持大规模的数据并行训练。DCN交换机使用Clos拓扑组织。DCN域可以扩展到10万个或更多的NPU。
4 架构的实现
在UB-Mesh-SuperPod中,4D-FullMesh + Clos的拓扑迫切需要高效的混合路由支持:除了通过LRS和HRS进行基于交换机的路由外,NPU本身也通过UB控制器具备路由能力。UB-Mesh中大量的路由器和分层拓扑给路由系统带来了新的复杂性。我们确定路由系统必须满足以下五个关键要求:
* 支持混合拓扑:路由系统必须满足UB-Mesh的4D-FullMesh + Clos混合拓扑所带来的要求。
* 高效转发:由于每个NPU也是一个路由器,且整个系统包含大量NPU,为了节省NPU硬件资源,路由系统必须高效地处理路由表查找和转发操作。
* 支持非最短路径:在nD-FullMesh拓扑中,两个端点之间存在多条距离不同的可能路径。系统应能够使用非最短路径,以最大化整个网络的带宽利用率。
* 快速故障恢复:为确保可靠性和可用性,路由系统必须能从任何故障中快速恢复,减少对训练过程的影响。
* 无死锁:最后,整个网络系统必须在没有死锁风险的情况下运行,确保数据流顺畅不中断。
然而,现有的用于Clos、Torus和DragonFly拓扑的路由技术,如表4所列的最长前缀匹配(LPM)、基于主机的路由和维度有序路由(DOR),都无法满足我们所有的要求。因此,我们提出了一种全路径路由(All-Path-Routing, APR)技术和通过直接通知的快速故障恢复技术来满足这些需求。
表4: 路由系统比较
4.1 全路径路由(APR)
UB-Mesh的全连接网格网络架构在任意两个NPU之间提供了多条路径。传统的路由策略,如最短路径优先(SPF),如图10-(a)所示,常常导致网络带宽利用不足,并且容易受到链路故障的影响。为了提高网络效率和弹性,我们在UB-Mesh中提出了全路径路由(APR)机制。如图10-(b)所示,APR利用了源和目标节点之间所有可用的路径。这种灵活性允许动态路径切换,以应对故障或拥塞,从而提高网络鲁棒性。为高效实现此目标,APR利用了三种底层机制:(1) 源路由(Source Routing, SR)机制以实现自适应路由。(2) 结构化寻址与线性表查找技术以最小化路由表查找和转发开销。(3) 拓扑感知的无死锁流控制以避免死锁。
图10: 最短路径路由 VS. 全路径路由
4.1.1 源路由(Source Routing)。为了充分利用APR提供的路径,一种实用的方法是利用源路由(SR)机制。如图11所示,一个8字节的SR头部被添加到原始数据包头部,其中包含转发指令。每个路由器使用一个4位的ptr
来指示12位位图字段中的一个位偏移量。第ptr
位的值指定了在第ptr
跳的转发方式(即,位=1表示在此跳进行SR转发,0表示传统转发)。在SR转发的情况下,位图字段也用于定位六个指令字段之一,该字段告诉如何转发此数据包。SR信息高度紧凑,因此开销很低。
图11: 源路由头部格式
4.1.2 结构化寻址与线性表查找。为了最小化每个NPU的UB IO控制器中的路由表查找和转发开销,APR路由系统利用了UB-Mesh拓扑的结构,并采用了一种结构化寻址和线性表查找机制。具体来说,地址空间根据网络元素的物理位置(如Pods、机柜和板卡)被划分为段。由于一个段内的NPU共享相同的前缀,因此只需要存储短的段地址,NPU可以通过相对于段地址的线性偏移来寻址。这种设计显著减少了表空间,加快了路由表的生成和分发,加速了收敛,并能快速响应状态变化,例如在故障恢复操作中。
4.1.3 拓扑感知的无死锁流控制。鉴于UB-Mesh的nD-FullMesh拓扑包含复杂的环形结构,并且APR机制也支持多跳路由,使用有限的虚拟通道(VL)资源实现无死锁流控制具有挑战性。为了克服这一挑战,我们提出了TFC(Topology-Aware Deadlock-Free Flow Control)算法。该算法在仅使用2个VL资源的情况下,最小化了VL资源的使用,同时实现了无死锁的全路径路由。TFC算法使用信道依赖图(CDG)对死锁进行建模,并将UB-Mesh的拓扑划分为子图。在每个子图内,拓扑转向规则和VL约束被统一为单个集合。它应用N维跨维度破环原则将该集合分解为单VL无环子集,然后应用同维度破环原则计算幂集元素的排列和笛卡尔积。这生成了所有路径组合和VL映射,确保了无死锁操作。由于篇幅限制,此处省略了详细解释。
4.2 通过直接通知的快速故障恢复
当网络中发生链路故障时,传统的路由系统通常采用逐跳通知。如图12左图所示,当链路1-3发生故障时,该信息由节点1和3逐跳传播,这通常会消耗很长的延迟。在UB-Mesh中,由于每个节点都有一组确定的通信目标,我们可以在链路故障时通过直接通知这些节点来加速路由收敛。如右图所示,当链路1-3发生故障时,该信息会根据预先计算的路由关系直接发送到节点6。通过这种拓扑感知的直接通知,可以大大减少控制平面的开销。
图12: 从逐跳通知到直接通知
5 性能最大化
尽管UB-Mesh的硬件架构设计与LLM训练的流量模式相匹配,但如果训练任务没有有效地分布在计算资源上,那么在这种分层AI集群上运行工作负载可能会导致利用率低下。此外,集合通信的性能对整体训练效率至关重要。为确保LLM训练期间的最佳系统性能,我们引入了几种拓扑感知的优化策略来进一步提升系统性能。
5.1 拓扑感知的集合通信
为了在UB-Mesh上优化集合通信并充分利用分层直接互连的带宽,我们提出了利用UB-Mesh的全路径路由(APR)能力的拓扑感知集合通信算法(见4.1节)。我们以All-Reduce和All-to-All为例来阐述我们的概念:
-
All-Reduce:我们提出了一种多环(Multi-Ring)算法,用于在UB-Mesh上高效地执行All-Reduce。我们首先通过统一的抽象对网络拓扑进行建模,考虑节点数、节点间连接、链路带宽和延迟等因素。接下来,我们使用逻辑多环算法将集合通信与路径映射相结合,确保路径的独占使用,避免流量冲突。如图13-(a)所示,“原始路径”代表默认映射。未包含在这些路径中的空闲链路,通过APR机制被利用来增强带宽。最后,如图13-(b)所示,我们优化了跨多条路径的流量划分,以减轻瓶颈,并通过APR最大化利用借用带宽带来的好处。
图13: 用于AllReduce的多环算法 -
All-to-All:在All-to-All通信场景中,我们考虑两个代表性的用例:
- 通用All-to-All:如图14-(a)所示(为简化起见,图中展示了一个2D网格,所提技术可扩展至更高维度),当UB-Mesh中的源节点向不同的目标节点发送不同数据时,我们采用多路径All2All优化。具体来说,每个元素(向量或张量)被分成两个分区,它们沿着X-FullMesh和Y-FullMesh互连同时传输,并最多使用一跳转发到达目的地。这种策略保证了UB-Mesh的nD-FullMesh架构中的高带宽利用率。
- 广播 + 归约:对于涉及令牌分发和专家数据收集的All-to-All操作【【7,DeepSeek-V3 Technical Report,2024,arXiv】】,其语义等同于重叠多个广播和归约操作。如图14-(b)和(c)所示,我们可以采用分层广播/归约来节省带宽使用,充分利用UB-Mesh的分层拓扑。
图14: 多路径和分层All-to-All
5.2 拓扑感知的并行化
给定一个训练任务和系统,我们确定最优的并行化策略,以充分利用UB-Mesh的高带宽局部互连。我们采用一种拓扑感知的并行化机制来优化LLM训练的模型和数据切分。如图15所示,我们的方法包括:
* 步骤①:生成可行的并行化配置,并将其映射到UB-Mesh架构上。
* 步骤②:使用拓扑感知的通信成本模型评估通信成本。
* 步骤③:迭代地最小化通信开销,以找到最优配置。
图15: 拓扑感知的并行化
该机制有两个重要要求:(1)构建一个合适的搜索空间以平衡效率和性能。(2)确保成本模型尽可能准确。对于要求(1),我们使用基于优先级的启发式方法来修剪搜索空间:通信量大的TP和SP(或CP)优先分配给高带宽域,而用于梯度更新的PP和DP优先级最低。对于需要EP的MoE模型,我们强制SP*DP是EP的整数倍。对于要求(2),我们精确地建模了APR和拓扑感知集合通信在UB-Mesh拓扑上的行为,并使用一个精确的内部仿真基础设施来校准该模型。
A4 实验环境
- 数据集/模型:基准测试工作负载包括:LLAMA-70B和GPT3-175B作为密集模型,以及采用MoE技术的稀疏模型GPT4-2T(参数根据【【23,GPT-4 Architecture & Infrastructure,2023,SemiAnalysis】】推断)。为了评估更大模型的性能,还包括了Dense-1T和MoE-10T模型。(见表5)
- 硬件配置:实验基于华为自研的硬件平台,包括NPU、CPU、低基数交换机(LRS)、高基数交换机(HRS)等组件。所有组件通过统一总线(UB)互连。实验规模从128 NPU扩展至最高64K NPU,SuperPod规模可达8K NPU。
- 软件配置:实验结果通过一个内部开发的仿真基础设施获得,该仿真器已经与真实的PoC(Proof of Concept)硬件进行了对齐,以评估集群规模的LLM训练性能。
表5: 基准测试模型
A4 实验结果
6.2 机柜内架构探索
- 实验内容:比较了四种不同的机柜内网络架构(如图16):(a) 2D-FM(UB-Mesh方案),(b) 1D-FM-A(板内直连,板间LRS),(c) 1D-FM-B(板内直连,板间HRS),以及(d) Clos(全交换机连接)。实验在8K NPU规模下进行,评估了不同序列长度下的训练吞吐量,并以Clos架构为基线进行归一化。
- 实验结果:
- 与Clos架构相比,2D-FM架构(UB-Mesh)在不同模型上的训练性能达到了93.2%至95.9%(图17-a)。
- 1D-FM-A和1D-FM-B的性能损失更小,但相比2D-FM的性能提升有限(大多在3%以内)。
- 在不同序列长度下,2D-FM的性能均与Clos基线保持在7%的差距内(图17-b)。
- 结论:UB-Mesh的2D-FM机柜内架构以远低于Clos架构的硬件成本,实现了与其相近的训练性能。
图16: 基线机柜内架构
图17: 不同机柜内拓扑的性能
6.3 机柜间架构探索
- 实验内容:比较了UB-Mesh的2D-FM机柜间直连架构(如图18-a)与基准Clos架构(如图18-b)。2D-FM架构还测试了三种路由策略:Shortest(仅最短路径)、Detour(使用APR的全路径)和Borrow(借用交换机带宽)。
- 实验结果:与理想的Clos架构相比,2D-FM的性能差距可以忽略不计。在使用Detour和Borrow策略后,对于对机柜间通信敏感的GPT4-2T模型,性能差距从0.73%缩小到仅0.46%(图19)。
- 结论:2D-FM机柜间互连架构的性能几乎与昂贵的Clos架构相同。
图18: 机柜间网络架构
图19: 不同机柜间互连的端到端性能比较
- 机柜间带宽探索:
- 实验内容:在8K SuperPod规模下,评估了不同机柜间带宽(每NPU UB x4, x8, x16, x32 IO)对性能的影响。
- 实验结果:对于短序列(8K-32K),最优带宽为UB x16;对于长序列(64K-10M),最优带宽为UB x32。带宽从x16提升到x32对长序列场景性能提升更显著(1.85%),因为部分TP和SP流量需要跨机柜(图20-b)。
- 结论:UB-Mesh默认每NPU分配UB x16 IO用于机柜间通信,以平衡成本和性能,并支持根据工作负载调整内外网带宽比例。
图20: 机柜间带宽探索
6.4 成本效益比较
- 实验内容:使用内部数据估算了不同架构的CapEx(资本支出)和OpEx(运营支出),并根据公式(见下图)计算成本效益。
- 实验结果:
- CapEx:UB-Mesh(4D-FM+Clos)架构相比其他基于Clos的变体,CapEx降低了1.18倍至2.46倍。网络基础设施成本在系统总成本中的占比从67%降至20%(图21)。与基准Clos相比,高基数交换机使用量减少98%,光模块减少93%。
- OpEx:由于交换机和光模块使用量大幅减少,UB-Mesh的OpEx相比Clos降低约35%。
- 总成本效益:综合性能和TCO(总拥有成本),UB-Mesh的成本效益比基准Clos高出2.04倍。
图21: CapEx比较
6.5 线性度评估
- 实验内容:根据公式(见下图)评估了UB-Mesh在不同集群规模下的性能线性扩展能力。基准规模因任务而异(如LLAMA2-70B为128 NPU)。
- 实验结果:在1倍至32倍的扩展规模下,所有任务的线性度均超过100%。当扩展到64倍规模(64K NPU)时,GPT4-2T和Dense-1T模型的线性度有所下降,但仍保持在95%以上(图22)。
- 结论:UB-Mesh表现出优异的线性扩展能力。
图22: 在序列长度256K下的线性度分析
6.6 网络可靠性分析
- 实验内容:根据公式(见下图)和各组件的年化故障率(AFR),估算了UB-Mesh和Clos架构的MTBF(平均无故障时间)和可用性。
- 实验结果:由于优先使用更可靠的电缆直连,UB-Mesh网络模块的AFR大大降低(表6)。UB-Mesh的MTBF达到98.5小时,是Clos的7.14倍。假设MTTR(平均修复时间)为75分钟,UB-Mesh的可用性为98.8%,相比Clos的91.6%提升了7.2%。
- 结论:UB-Mesh架构显著提升了网络可靠性和系统可用性。通过自研网络监控工具将MTTR缩短至13分钟内,可用性可进一步提升至99.78%。
表6: MTBF估算
A5 结论
本文提出了UB-Mesh,一个专为下一代LLM训练设计的新型数据中心网络架构。UB-Mesh采用nD-FullMesh网络拓扑,以节省交换机和光模块的使用,同时适应LLM工作负载的分层局部化流量模式。我们仔细考虑了架构、物理实现和网络系统优化,并提出了多种技术来克服各种挑战。与传统的Clos网络相比,UB-Mesh在提供相近LLM训练性能的同时,系统级成本效益提高了2.04倍,可用性提高了7.2%。此外,UB-Mesh在多个LLM训练任务上实现了超过95%的线性度。
A6 附录(讨论)
用于集合通信的协处理器。值得一提的是,在UB IO控制器中还有一个名为集合通信单元(Collective Communication Unit, CCU)的特殊协处理器,用于卸载集合通信任务。具体来说,CCU执行指令,主动从HBM读写数据,发起NPU间传输,并利用片上SRAM缓冲区执行在线数据归约。它消除了从应用程序内存缓冲区到通信缓冲区的冗余数据拷贝。这种设计有效地减轻了HBM带宽消耗,并通过基于校验位的细粒度同步机制保持了确定性的归约顺序。CCU还可以与计算核心无缝协作,以实现高效的计算-通信重叠。
对大规模专家模型的支持。除了密集的和常规的MoE模型,我们注意到人们正在积极探索包含海量专家的MoE模型【【6,DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model,2024,arXiv】,【7,DeepSeek-V3 Technical Report,2024,arXiv】】。这种设计通常需要大规模、细粒度的all-to-all通信。UB-Mesh能够通过多路径和分层的all-to-all优化(见5.1节)以及基于UB的Load/Store数据传输来高效支持此类海量专家模型。CCU还可以高效地卸载all-to-all操作,以节省宝贵的计算核心资源,正如在【【7,DeepSeek-V3 Technical Report,2024,arXiv】】中所预期的那样。
💬 评论讨论
欢迎在这里分享您的想法和见解!