UB-Mesh:一種分層局部化的n維全互連數據中心網絡架構
- ?標題: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
- 鏈接:https://arxiv.org/abs/2503.20377v1
- 時間:26 Mar 2025
摘要
隨著大規模語言模型(LLM)的規模持續擴大,其所需的計算能力和帶寬需求也在不斷提升。為應對這一挑戰,我們提出了UB-Mesh,一種新穎的AI數據中心網絡架構,旨在提升可擴展性、性能、成本效益以及可用性。與傳統數據中心提供對稱的節點間帶寬不同,UB-Mesh采用了一種分層局部化的n維全互連(nD-FullMesh)網絡拓撲結構。該設計充分利用了LLM訓練中的數據局部性,優先采用短距離直接互連,以最小化數據移動距離并減少交換機的使用。
盡管UB-Mesh的n維全互連拓撲在理論上具有多項優勢,但其具體的架構設計、物理實現以及網絡系統優化帶來了新的挑戰。為實現UB-Mesh的實際構建,我們首先設計了基于4維全互連拓撲的UB-Mesh-Pod架構。UB-Mesh-Pod通過一系列硬件組件實現,這些組件構成了基礎構建模塊,包括專門設計的NPU、CPU、低基數交換機(LRS)、高基數交換機(HRS)、網卡(NIC)等。這些組件通過一種新穎的統一總線(Unified Bus, UB)技術實現互連,該技術支持靈活的IO帶寬分配和硬件資源池化。在網絡系統優化方面,我們提出了名為全路徑路由(All-Path-Routing, APR)的先進路由機制,以高效管理數據流量。這些優化措施結合拓撲感知的性能增強和諸如64+1備份設計的穩健可靠性機制,使得UB-Mesh相比傳統的Clos架構實現了2.04倍的成本效益提升、7.2%的網絡可用性提升,并在多種LLM訓練任務中實現了95%以上的線性擴展能力。。
1.引言
新興的大規模語言模型(LLM)[1, 4, 6, 8, 10, 18, 22, 26, 27, 31, 32]正在改變AI產業和人類社會。根據規模定律(Scaling Laws)[5, 13],LLM通過增加模型參數和訓練數據量持續提升其理解、生成和推理能力。然而,這一趨勢對底層訓練系統和基礎設施提出了日益嚴峻的挑戰,迫使下一代AI數據中心滿足以下要求:
圖1:傳統Clos數據中心架構與UB-Mesh的比較
R1:大規模。隨著模型規模和訓練數據量的增加,需要越來越多的NPU(神經網絡處理單元)或GPU在合理的時間內完成訓練。例如,Llama-3的預訓練使用1.6萬個GPU需要54天[8]。領先公司最近的公告顯示,已成功部署了擁有10萬GPU的AI訓練系統[30]。可擴展的基礎設施對于支持LLM技術的持續發展至關重要。
R2:高帶寬。在LLM訓練系統中,AI計算節點(NPU/GPU)需要的互連帶寬超過每節點3.2 Tbps[12, 17],大約是當代數據中心中典型CPU節點互連帶寬的10倍。因此,最先進的AI訓練系統的總帶寬比當前基于CPU的基礎設施即服務(IaaS)系統高出10到100倍。
R3:成本效益。構建大規模AI數據中心需要巨額硬件投資,資本支出(CapEx)往往達到數十億美元。為實現所需的10到100倍總互連帶寬增長,若使用傳統的對稱Clos數據中心網絡架構(見圖1-(a)),互連成本也將增加10到100倍。優化網絡基礎設施為提升成本效益提供了重要機會。此外,降低運營成本(OpEx),包括能耗和維護費用,對于確保整體成本效益同樣至關重要。
R4:高可用性。擁有10萬計算節點和約100萬個光模塊的大型LLM訓練集群面臨顯著的可用性挑戰。當前統計數據顯示,即使每條鏈路的平均故障間隔時間(MTBF)為5年,整個10萬GPU的AI集群的原始MTBF仍下降至不到30分鐘[24]。為此,網絡架構設計不僅需提升硬件可靠性,還需融入容錯機制,以應對互連、計算資源、控制系統和存儲的故障。
同時實現這些目標極具挑戰性,需要向先進的下一代數據中心網絡架構設計轉變。我們認為,設計下一代AI數據中心應以以下三個原則為核心:
P1:流量模式驅動的網絡拓撲。與傳統數據中心工作負載通常產生均勻且隨機的流量不同,大規模語言模型(LLM)的訓練流量具有確定性并表現出強烈的數據局部性。例如,張量并行(Tensor Parallelism)[25]操作所需的集合通信通常占總數據流量的50%以上,且通常發生在8-64個相鄰NPU的集群內。相比之下,數據并行(Data Parallelism)產生的集合通信占總流量的不到2%,但通常需要長距離傳輸。因此,層次化、局部化的網絡架構對于適配這些流量模式至關重要。
P2:拓撲感知的計算與通信。在分層局部化的數據中心網絡中,有效運行LLM訓練是另一重大挑戰。如果訓練任務未在計算資源上優化分布,或網絡系統未得到充分優化,AI集群可能因流量擁塞或帶寬利用不足而性能低下。為解決這一問題,并行策略選擇、路由、集合通信、負載均衡等因素必須與網絡拓撲精確對齊。
P3:容錯自愈系統。LLM訓練系統必須具備自愈能力以確保魯棒性。鏈路故障時,路由系統應自動切換至備用路徑;同樣,若NPU發生故障,應有機制無縫激活備用NPU,以維持系統完整性和LLM訓練的連續性。
為滿足要求R1-R4并遵循原則P1-P3,我們提出了創新的UB-Mesh架構。如圖1-(b)所示,UB-Mesh采用n維全互連(nD-FullMesh)網絡拓撲,通過遞歸構建全互連拓撲——從板上相鄰NPU間的1維連接開始,擴展到機架內相鄰1維網格間的2維連接,再通過互連更廣范圍的相鄰高維網格擴展至3維及以上。這種分層局部化的網絡架構最小化傳輸跳數并優化每跳距離,優先采用直接互連而非間接長距離交換,從而減少對交換機和光模塊的依賴,滿足R1和R3。此外,它根據LLM訓練的傳輸需求實現層次化帶寬分配,為短距離通信提供高帶寬,為長距離通信提供較低帶寬,滿足R2和R3并遵循P1。
遵循P2,我們深入研究了先進的網絡和系統優化機制以增強UB-Mesh架構。具體而言,我們引入了全路徑路由(All-Path Routing, APR)技術,以充分利用直接連接鏈路的帶寬。APR結合源路由、結構化尋址與線性表查找以及無死鎖流量控制機制,實現自適應路由、最小化轉發開銷并避免死鎖。此外,我們加入了拓撲感知的快速故障恢復機制以提升可靠性。為進一步優化性能,我們提出了拓撲感知的集合通信和并行優化,以提升訓練期間的帶寬利用率。
為滿足R4并遵循P3,UB-Mesh采用64+1高可用性設計:每個機架包含一個額外的備用NPU。當系統中NPU發生意外故障時,備用NPU被激活以恢復功能,確保LLM訓練任務不間斷繼續。此外,路由系統通過一種新穎的直接通知技術在鏈路故障時實現快速恢復。
我們精心設計了UB-Mesh的硬件和系統堆棧,綜合考慮了各種工程約束和權衡。具體實現中,UB-Mesh-Pod采用4維全互連拓撲,使UB-Mesh可無縫擴展至8000個NPU,形成支持下一代AI數據中心構建的高帶寬域。為此,我們開發了一系列硬件組件作為基礎構建模塊,包括NPU、CPU、低基數交換機(LRS)、高基數交換機(HRS)、網卡(NIC)等。
與采用多種互連技術(如PCIe、NVLINK、IB和RoCE)的基線系統不同,UB-Mesh使用新穎的統一總線(Unified Bus, UB)技術實現所有組件互連。這種統一方法提升了IO資源分配的靈活性,其點對點通信能力支持高效的硬件資源池化,UB還為無縫跨層優化提供了機會。
全面評估表明,與非超訂閱的Clos網絡相比,UB-Mesh將高基數交換機使用量減少98%,光模塊使用量減少93%,實現系統級成本效益提升2.04倍。在多個LLM訓練任務中的實驗還顯示,與成本高昂的Clos網絡相比,UB-Mesh的性能降幅僅在7%以內。這種低成本與高性能的結合不僅滿足當前LLM訓練需求,還使其能夠有效應對未來的可擴展性挑戰。
2.背景與動機
2.1 LLM訓練中的“通信墻”
大規模語言模型(LLM)訓練是有史以來規模最大、計算和通信需求最為密集的并行計算應用[11, 18, 20–22, 25, 27, 32]。遵循所謂的規模定律(Scaling Laws)[3, 5, 13],LLM通過增加模型參數和訓練數據量來提升性能。因此,LLM需要越來越多的AI加速器以在合理的時間內完成訓練。例如,開源的Llama-3.1模型在1.6萬個GPU上進行訓練[8],而下一代LLM模型已使用10萬GPU進行訓練[24, 29]。
標準的訓練過程包括重復的訓練迭代,每個迭代包含前向傳播以計算損失、反向傳播以確定梯度,以及優化器步驟以調整模型參數。為充分利用分布式計算能力,LLM訓練通過多種并行策略將數據、模型和激活值分割到數萬個NPU上。在每個迭代中,NPU之間頻繁交換數據以分發輸入數據、同步激活值和梯度等。隨著訓練系統規模的擴大,數據移動成為系統中成本最高的部分[9, 19]。若缺乏強大的NPU間通信能力,訓練過程很容易受限于“通信墻”(Communication Wall)。
2.2 LLM訓練中數據流量的局部性
圖2:LLM訓練中的并行性
如圖2所示,LLM訓練通常涉及多種并行技術,以下逐一介紹:
張量并行(Tensor Parallelism, TP):TP以行或列的方式分割模型層,并將子層放置在多個NPU上并行計算[25]。其主要涉及AllReduce操作以合并分布式的部分結果。
序列并行(Sequence Parallelism, SP):SP(在某些論文中也稱為上下文并行)通常用于將序列分割到多個NPU以實現并行處理。SP依賴RingAttention[15]技術,并可采用AllGather操作收集不同NPU的部分結果。
專家并行(Expert Parallelism, EP):對于采用專家混合(Mix-of-Experts, MoE)技術的LLM模型,密集MLP層被MoE層替代,每個MoE層包含若干“專家”。執行時專家被稀疏激活。EP將專家分布在不同NPU上,輸入標記通過All2All通信動態發送至目標專家。
流水線并行(Pipeline Parallelism, PP):與TP分割每個模型層不同,PP將層分布到多個設備上,并以流水線方式執行前向和反向傳播。PP涉及低開銷的P2P通信以傳輸層間激活值,但需要高效的調度算法[11, 16]以最小化流水線中的空泡。
數據并行(Data Parallelism, DP):DP在多個NPU上復制模型和優化器狀態,每個副本并行處理一部分輸入批次。訓練過程中需通過AllReduce操作同步梯度。
表1:LLM訓練中數據流量的分析
這些并行技術共同將訓練任務分布到數千個NPU上。需要注意的是,并非所有并行技術產生的數據流量均等。根據我們基于內部MoE-2T模型的分析(如表1詳述),通信強度呈層次化并表現出強烈的局部性。具體而言,TP和SP約占總流量的97%,而其余并行技術通常產生不到2%的流量。其他模型架構可能表現出略有不同的數據流量分布,但同樣具有較強的局部性[28]。因此,架構設計應優先采用層次化的網絡帶寬分配方式。
2.3 數據中心網絡架構
圖3:傳統數據中心網絡架構
設計大規模LLM訓練集群涉及一個關鍵問題:如何將眾多NPU資源組織成一個協調的分布式系統?這需要確定合適的數據中心網絡架構。傳統數據中心和超算系統探索了多種拓撲結構,但這些拓撲未必適用于大規模LLM訓練。圖3展示了一些常見拓撲:
- Clos:傳統數據中心通常采用Clos架構,通過兩到三層交換機對稱連接NPU/CPU,提供高性能并適應多種流量模式。然而,其高成本是一個缺點,主要源于對高性能交換機和光模塊的廣泛使用。
- 3D Torus:與依賴交換機互連的Clos不同,3D Torus直接連接相鄰NPU。雖然降低了成本,但3D Torus的NPU間帶寬較低,且難以適應復雜的通信模式,例如MoE模型中常見的All-to-All通信[20]。
- Dragon-Fly:Dragon-Fly拓撲[14]將交換機組織為互連組,每組內交換機彼此直接相連,組間通過長距離電纜連接,數據包通過間接路由傳輸。該架構在傳統數據中心和高性能計算(HPC)環境中降低了維度。Dragon-Fly比Clos成本更低,但因需滿足完整的NPU-交換機帶寬需求仍較昂貴,且其架構在LLM訓練流量(特別是P2P和AllReduce場景)中表現不佳。
- Fugaku Tofu:Fugaku Tofu[2]提出了一種獨特的6D Torus拓撲用于HPC應用。然而,與3D Torus類似,其NPU間帶寬較低,且難以適應復雜的通信模式。
- 綜上所述,如何為大規模LLM訓練設計一個高性能且成本效益高的數據中心網絡架構仍是一個開放性問題。
3.UB-Mesh架構
為滿足第1節中提出的要求R1-R4并遵循原則P1-P3,我們提出了UB-Mesh架構。UB-Mesh采用了一種新穎的n維全互連(nD-FullMesh)拓撲結構,最大限度地利用直接電線互連,從而減少對昂貴高帶寬交換機和光模塊的需求。接下來,我們將介紹一系列專為構建UB-Mesh設計的硬件模塊。最后,我們將詳細討論UB-Mesh-Pod的架構設計,即UB-Mesh的具體4維全互連實現,同時充分考慮了若干工程約束和權衡。
3.1 nD-FullMesh拓撲
圖4:UB-Mesh的n維全互連拓撲及可能的物理映射
如圖4所示,UB-Mesh引入了一種新穎的n維全互連拓撲,強調跨所有網絡層級的直接互連。該拓撲從相鄰節點形成1維全互連開始,其中同一層級的每個節點與該層內其他所有節點相連。基于這一概念,相鄰1維全互連的節點之間再進行直接互連,形成2維全互連。這一過程遞歸進行,相鄰的2維全互連形成3維全互連,依此類推,最終根據需要擴展至n維全互連。這種靈活的虛擬拓撲可以無縫映射到各種物理NPU組織形式:單板上的1維全互連、機架內的2維全互連、機架行間的3維全互連、樓層機架組內的4維全互連,甚至跨越整棟建筑的5維全互連及更高維度。
圖5:高維拓撲與靈活的帶寬分配
這種分層局部化的拓撲在構建下一代LLM訓練數據中心時具有多項優勢。首先,由于nD-FullMesh在每個網絡維度上形成緊密耦合的直接連接域,并可提供逐層超訂閱的帶寬,它充分利用了LLM訓練固有的數據局部性和從密集到稀疏的流量模式(見第2.2節)。此外,我們可以靈活調整每個維度的每節點帶寬分配,以滿足未來LLM訓練任務的具體需求。例如,如圖5-(a)所示,假設存在6維全互連拓撲,我們可以通過在NPU的IO模塊中分配不同數量的互連鏈路來調整??????和??????維度的帶寬,如圖5-(b)所示。
其次,與需要大量交換的Clos架構不同,這種直接連接的拓撲在減少傳輸距離方面具有巨大潛力。結合數據放置和集合通信優化(見第5.2節),大多數傳輸(通常涉及TP或SP,如表1所示)可在0-2跳內完成,大幅降低數據移動開銷。
表2:不同類型鏈路的使用估算
第三,與Clos等其他流行拓撲相比,nD-FullMesh拓撲減輕了對高帶寬交換機和長距離光互連的依賴。相反,它最大限度地利用短距離直接互連。例如,根據表2的估算,短距離無源電線占據了總電纜消耗的86.7%。這一設計顯著降低了交換機和光模塊的成本,同時提升了系統可靠性,因為電線和連接器比光模塊更穩定。
最后,與不適合復雜集合通信操作(如All-to-All)的3D-Torus拓撲不同,全互連拓撲能夠高效支持這些操作,具體將在第5節中討論。
3.2 UB-Mesh的構建模塊
所提出的nD-FullMesh拓撲在理論上具有多項優勢,但其實踐實現和系統性能的提升帶來了新的挑戰,需要全面考慮各種工程約束和權衡。為此,我們設計并制造了一系列硬件模塊,作為UB-Mesh的基礎構建塊。這些模塊通過一種名為統一總線(Unified Bus, UB)的新型互連機制實現連接。具體細節如下:
3.2.1 硬件模塊
表3:UB-Mesh的主要構建模塊
表3列舉了UB-Mesh的主要組件。其核心是作為AI計算單元的NPU。NPU配備兩個UB IO控制器,總計提供x72條UB通道。負責執行主機程序的CPU也配備一個UB IO控制器,提供x32的UB IO。
盡管UB-Mesh優先采用直接互連,但交換能力對于滿足特定需求仍不可或缺。例如,在圖5中聚合??和??????維度的NPU IO有助于實現高效的機架間互連。此外,我們期望CPU-NPU比例及其綁定關系能夠靈活調整,以實現NPU和內存資源的高效池化,這可以通過交換實現。
為在UB-Mesh-Pod內提供成本效益高的交換能力,我們開發了輕量級低基數交換機(LRS)。其制造成本低,可高效聚合機架間IO帶寬并支持CPU-NPU通信。
盡管UB-Mesh僅使用低基數交換機即可實現,但實際網絡約束通常將其規模限制在一定范圍內(UB-Mesh-Pod)。因此,UB-Mesh還包括高基數交換機(HRS),提供x512的UB IO容量以支持Pod級交換。UB-Mesh的詳細架構設計和考慮將在第3.3節中介紹。
3.2.2 統一總線(UB)互連
圖6:統一總線互連與傳統混合互連的對比
連接這些異構硬件模塊需要特定的互連技術。如圖6-(a)所示,傳統的基于GPU的LLM系統通常采用多種互連方式,而UB-Mesh通過統一總線(UB)技術實現所有組件的互連,其優勢包括:
- 靈活的IO資源分配:UB互連與特定用例解耦,支持芯片內不同類型IO的靈活資源分配,如圖5所示。NPU間帶寬和CPU-NPU帶寬也可根據具體需求靈活調整,因為它們使用相同的UB鏈路。
- 硬件資源池化:UB的點對點通信能力支持包括DDR DRAM、CPU、NPU和NIC在內的高效硬件資源池化。例如,在圖6-(b)中,CPU和NPU通過UB互連實現池化,以提升資源利用率。
- 系統優化收益:統一互連消除了協議轉換的需要,大幅降低開銷并簡化驅動程序、通信庫和框架的設計與優化。
3.3 架構概覽
圖7:UB-Mesh的架構設計
通過組合表3中列出的硬件構建模塊并通過統一總線(Unified Bus, UB)進行互連,我們提出了UB-Mesh的整體架構設計。如圖7-(a)所示,我們按照所提出的nD-FullMesh拓撲實現了一個UB-Mesh-Pod。具體而言,在該架構中,我們在每個機架內創建了一個2D全互連拓撲,并將其擴展到機架之外的另一個2D全互連拓撲,從而形成一個4D全互連拓撲。我們選擇不在這一代產品中將UB-Mesh擴展至5D全互連或更高維度,以在成本效益和靈活性之間實現工程平衡,在4D全互連之外仍采用Clos拓撲。以下為詳細介紹:
3.3.1 在機架內實現2D全互連
圖8:UB-Mesh機架的硬件實現
如圖7-(b)所示,每個機架配置為2D全互連拓撲,集成了多個組件以確保高效通信和資源利用。機架的核心由八個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),以降低成本,同時確保設備間的無阻塞通信。總體而言,這些背板交換機輸出四個UB x256 IO。
3.3.2 64+1高可用性設計
正如第1節所述,在大規模AI數據中心中確保系統可用性是一個關鍵挑戰。系統中單個NPU可能會發生意外故障。雖然軟件系統可以檢測這些故障并使用剩余健康NPU重啟訓練,但由于計算能力和系統帶寬的減少,系統性能將受到顯著影響。
圖9:通過啟用備用NPU實現容錯
為提升系統可用性,我們引入了特殊的64+1設計:如圖7-(b)和圖8-(b)所示,系統包含64個常規NPU以及一個額外的備用NPU。該備用NPU通過LRS與64個常規NPU相連。當任何NPU發生故障時,備用NPU將被策略性地啟用。如圖9所示,當NPU-3發生故障時,管理系統激活備用NPU(圖中節點B)以替換NPU-3。與NPU-3相關的原始直接連接鏈路將被重定向。例如,路徑5-3被重定向為路徑5-LRS-B。盡管此策略將原始直接連接變為單跳路由,略微增加了傳輸延遲,但相比簡單屏蔽NPU-3并在剩余七個NPU上運行任務,其效果要優越得多。在機架內帶寬為主要關注點的LLM訓練場景中,增加的延遲對整體訓練性能影響微乎其微。
3.3.3 在UB-Mesh-Pod中擴展至4D全互連
在機架內,我們將64個NPU組織為2D全互連拓撲。在機架之外,我們實現另一個2D全互連來組織機架,形成名為UB-Mesh-Pod的4D全互連架構。
如圖7-(a)和(c)所示,一行中四個相鄰機架被組織為緊密耦合的1D全互連,機架通過每個機架中的LRS端口直接互連。在每個機架列中,機架也直接相連,形成一個16機架的2D全互連拓撲。此配置中的每條鏈路代表一個UB x128 IO,如圖8-(d)所示。我們沿兩個維度連接四個相鄰機架以構建機架間全互連,這是考慮到有源電線的覆蓋范圍得出的最佳點。由于每個機架有64個NPU,每個Pod有16個機架,一個4D全互連的UB-Mesh-Pod總共包含1024個NPU。
3.3.4 UB-Mesh-SuperPod及更高規模
基于1000規模的UB-Mesh-Pod,我們進一步構建了UB-Mesh-SuperPod,可容納多個UB-Mesh-Pod。考慮到當前云場景中小規模或中規模LLM訓練負載可能不會消耗整個SuperPod,我們選擇在Pod級互連中采用對稱的Clos拓撲,而非繼續使用全互連設計。這種設計允許云管理者根據用戶需求靈活劃分SuperPod,并保證每個劃分域內的全帶寬。如圖7-(c)所示,我們使用高基數Pod交換機(HRS)連接SuperPod中的每個機架,可擴展至8000個NPU。
最后,SuperPod中的機架還通過UB交換機(方案(a))或位于CPU板上的NIC(方案(b))連接到大規模數據中心網絡(DCN)。DCN域通常支持大規模數據并行訓練。DCN交換機采用Clos拓撲組織。DCN域可擴展至10萬NPU或更多。
4.實現架構
在UB-Mesh-SuperPod中,4D全互連與Clos混合拓撲迫切需要高效的混合路由支持:除了通過LRS和HRS進行的基于交換機的路由外,NPU本身通過UB控制器也具備路由能力。UB-Mesh中大量路由器和層次化拓撲為路由系統引入了新的復雜性。我們認為,路由系統必須滿足以下五個關鍵要求:
- 支持混合拓撲:路由系統必須滿足UB-Mesh的4D全互連與Clos混合拓撲所帶來的需求。
- 高效轉發:由于每個NPU同時也是路由器,且整個系統包含大量NPU,為節省NPU硬件資源,路由系統必須高效處理路由表查詢和轉發操作。
- 支持非最短路徑:在nD全互連拓撲中,兩個端點之間存在多個路徑,距離各異。系統應支持使用非最短路徑,以最大化網絡帶寬利用率。
- 快速故障恢復:為確保可靠性和可用性,路由系統必須快速從故障中恢復,減少對訓練過程的影響。
- 無死鎖:最后,整個網絡系統必須在無死鎖風險下運行,確保數據流暢通無阻。
表4:路由系統比較
然而,表4中列出的現有路由技術,如Clos、Torus和DragonFly拓撲上使用的長前綴匹配(LPM)、基于主機的路由和維度順序路由(DOR),均無法完全滿足我們的需求。因此,我們提出了全路徑路由(All-Path-Routing, APR)技術和基于直接通知的快速故障恢復技術,以滿足上述要求。
4.1 全路徑路由(APR)
圖10:最短路徑路由與全路徑路由
UB-Mesh的全互連網絡架構為任意兩個NPU之間提供了多條路徑。傳統的路由策略,如最短路徑優先(SPF)路由(見圖10-(a)),往往無法充分利用網絡帶寬,且易受鏈路故障影響。為提升網絡效率和魯棒性,我們在UB-Mesh中提出了全路徑路由(APR)機制。
如圖10-(b)所示,APR利用源節點與目標節點之間的所有可用路徑。這種靈活性允許動態切換路徑,以應對故障或擁塞,從而提升網絡魯棒性。為高效實現這一目標,APR依賴以下三種底層機制:(1) 源路由(SR)機制實現自適應路由;(2) 結構化尋址與線性表查詢技術,減少路由表查詢和轉發開銷;(3) 拓撲感知無死鎖流量控制,避免死鎖。
4.1.1 源路由
圖11:源路由頭部格式
為充分利用APR提供的路徑,一種實用方法是采用源路由(SR)機制。如圖11所示,在原始數據包頭部添加一個8字節的SR頭部,包含轉發指令。每個路由器使用一個4位指針(ptr)指示12位位圖字段中的位偏移。第??位的值指定第??跳的轉發方式(即,位值為1表示該跳采用SR轉發,0表示傳統轉發)。在SR轉發的情況下,位圖字段還用于定位六個指令字段之一,指示如何轉發該數據包。SR信息高度壓縮,因此開銷較低。
4.1.2 結構化尋址與線性表查詢
為減少每個NPU的UB IO控制器中的路由表查詢和轉發開銷,APR路由系統利用UB-Mesh拓撲的結構,采用結構化尋址和線性表查詢機制。具體而言,地址空間根據網絡元素的物理位置(如Pod、機架和板)劃分為多個段。由于同一段內的NPU共享相同前綴,僅需存儲短段地址,NPU可通過相對于段地址的線性偏移進行尋址。此設計顯著減少表空間,加速路由表生成和分發,提升收斂速度,并支持快速響應狀態變化,例如故障恢復操作。
4.1.3 拓撲感知無死鎖流量控制
鑒于UB-Mesh的nD全互連拓撲包含復雜的環結構,且APR機制支持多跳路由,在有限的虛擬通道(VL)資源下實現無死鎖流量控制具有挑戰性。為應對這一挑戰,我們提出了TFC(拓撲感知無死鎖流量控制)算法。該算法在僅使用2個VL資源的情況下,最小化VL資源使用,同時實現無死鎖的全路徑路由。
TFC算法通過通道依賴圖(CDG)建模死鎖,并將UB-Mesh拓撲劃分為子圖。在每個子圖內,拓撲導向規則和VL約束被統一為單一集合。算法應用N維跨維度斷環原則,將集合分解為單VL無環子集,隨后采用同維度斷環原則計算冪集元素的排列和笛卡爾積,生成所有路徑組合和VL映射,確保無死鎖運行。由于篇幅限制,詳細說明在此省略。
4.2 基于直接通知的快速故障恢復
圖12:從逐跳通知到直接通知
當網絡中發生鏈路故障時,傳統路由系統通常采用逐跳通知方式。如圖12左側所示,當鏈路1-3發生故障時,該信息由節點1和3逐跳傳播,通常耗時較長。在UB-Mesh中,由于每個節點具有確定的通信目標集,我們可以通過直接通知受影響的節點加速路由收斂。如圖12右側所示,當鏈路1-3發生故障時,根據預計算的路由關系,該信息直接發送至節點6。通過這種拓撲感知的直接通知方式,可大幅減少控制平面的開銷。
5.性能最大化
盡管UB-Mesh的硬件架構設計與LLM訓練的流量模式相匹配,但在這種層次化的AI集群上運行工作負載時,如果訓練任務未能在計算資源上有效分布,可能會導致低利用率。此外,集合通信的性能對整體訓練效率至關重要。為確保LLM訓練期間的系統性能最優,我們引入了若干拓撲感知優化策略,以進一步提升系統性能。
5.1 拓撲感知集合通信
為優化UB-Mesh上的集合通信并充分利用層次化直接互連的帶寬,我們提出了利用UB-Mesh全路徑路由能力(見第4.1節)的拓撲感知集合通信算法。我們以All-Reduce和All-to-All為例闡述我們的概念:
圖13:AllReduce的多環算法
All-Reduce:我們提出了一個多環(Multi-Ring)算法,以在UB-Mesh上高效實現All-Reduce。我們首先通過統一的抽象模型對網絡拓撲進行建模,考慮節點數量、節點間連接、鏈路帶寬和延遲等因素。接下來,我們將集合通信與路徑映射相結合,采用邏輯多環算法,確保路徑獨占使用,避免流量沖突。如圖13-(a)所示,原始路徑表示默認映射。未包含在這些路徑中的空閑鏈路通過APR機制被利用,以提升帶寬。最后,如圖13-(b)所示,我們優化了多路徑上的流量分區,以緩解瓶頸并通過APR最大化借用帶寬的收益。
All-to-All:在All-to-All通信場景中,我們考慮了兩種代表性用例:
圖14:多路徑與層次化All-to-All
(1) 通用All-to-All:如圖14-(a)所示,當UB-Mesh中的源節點(為簡化起見,圖中展示2D網格,提出的技術可擴展至更高維度)向不同目標節點發送不同數據時,我們采用多路徑All2All優化。具體而言,每個元素(向量或張量)被分為兩個分區,分別沿X-FullMesh和Y-FullMesh互連同時傳輸,并最多使用一跳轉發到達目標。這種策略保證了UB-Mesh nD全互連架構中的高帶寬利用率。
(2) 廣播+歸約(Broadcast + Reduce):對于涉及標記分發和專家數據收集的All-to-All操作[7],其語義等價于多個廣播和歸約操作的重疊。如圖14-(b)和(c)所示,我們可采用層次化廣播/歸約以節省帶寬使用,充分利用UB-Mesh的層次拓撲。
5.2 拓撲感知并行化
圖15:拓撲感知并行化
給定訓練任務和系統,我們確定最優并行策略,以充分利用UB-Mesh的高帶寬局部互連。我們采用拓撲感知并行化機制優化LLM訓練的模型和數據分割。如圖15所示,我們的方法包括:
步驟①:生成可行的并行配置并將其映射到UB-Mesh架構上。
步驟②:利用拓撲感知通信成本模型評估通信成本。
步驟③:迭代最小化通信開銷以找到最優配置。
該機制有兩個重要要求:(1) 構建適當的搜索空間以平衡效率和性能;(2) 確保成本模型盡可能準確。對于要求(1),我們使用基于優先級的啟發式方法剪枝搜索空間:涉及高通信量的TP和SP(或CP)優先分配到高帶寬域,而用于梯度更新的PP和DP優先級最低。對于需要EP的MoE模型,我們強制SP*DP為EP的整數倍。對于要求(2),我們在UB-Mesh拓撲上準確建模APR和拓撲感知集合通信的行為,并使用內部高精度仿真基礎設施校準模型。
6.評估
在本節中,我們探索了UB-Mesh的架構設計空間,并分析其相較于基線Clos架構的優勢。為便于更詳細的比較,我們將分析分為兩個層級:機架內架構比較和機架間架構比較。
6.1 實驗設置
表5:基準模型
表5列出了基準工作負載:Llama-70B和GPT3-175B為密集模型,而GPT4-2T為采用MoE技術的稀疏模型[20]。需要注意的是,GPT4的架構尚未正式發布,我們采用了推測參數[23]。為評估系統在更大模型上的性能,我們還包括了Dense-1T和MoE-10T模型。為探索UB-Mesh的架構,我們構建了一個內部仿真基礎設施,并與真實PoC硬件對齊,以評估集群規模的LLM訓練性能。
6.2 機架內架構探索
圖16:基線機架內架構
我們比較了不同的機架內網絡架構,如圖16所示:
(a) 2D-FM:UB-Mesh的架構。64個NPU通過電線直接互連,形成2D全互連網絡拓撲(X-Fullmesh + Y-Fullmesh)。該設計成本低,采用48個LRS(見表3)聚合機架間帶寬并實現CPU-NPU互連(為簡化起見,圖中省略了CPU)。
(b) 1D-FM-A:此替代架構保留1D X-Fullmesh,即每個板上的8個NPU仍直接互連。然而,跨板通信通過36個LRS實現。每個NPU有一個UB x16 IO連接到LRS,另一個UB x16 IO連接到四個高基數交換機(HRS)用于機架間通信。
(c) 1D-FM-B:此架構進一步用HRS替換LRS。每個背板中的四個LRS用于NPU到CPU通信。跨板NPU通信通過四個背板中的八個交換機實現。這些交換機還連接到機架間網絡,為每個NPU提供UB x32 IO用于機架間通信。
(d) Clos:此架構不使用NPU間直接連接,而是將64個NPU的所有端口連接到72個LRS,形成對稱的Clos拓撲。該架構提供最高靈活性,但需要大量交換資源。
圖17:不同機架內拓撲的性能
性能比較:在圖17中,我們固定機架間架構(2D-FM,見第6.3節),并比較不同機架內架構的訓練吞吐量,相對于Clos基線。SuperPod規模為8000(128個機架),評估的序列長度范圍為8000至1000萬。我們計算了不同序列長度的平均性能。
如圖17-(a)所示,與Clos架構相比,2D-FM架構實現了93.2%至95.9%的訓練性能。1D-FM-A架構顯示出較低的性能下降,對于Llama2-70B相比2D-FM提升了2.44%的性能。對于其他參數更多的模型,提升幅度小于1.6%。1D-FM-B架構由于更高的機架間帶寬,相比2D-FM顯示出略高的性能提升,超過3%,但提升幅度仍較小。
圖17-(b)探索了不同序列長度下的性能,相對于Clos架構的基線場景(所有模型平均值)。對于8000至3.2萬的序列長度,2D-FM架構實現了95.5%的性能,略低于1D-FM-A(98.1%)和1D-FM-B(99.2%)。對于6.4萬至1000萬的序列長度,2D-FM架構相比Clos架構實現了95.0%的性能。
可以看出,與Clos相比,2D-FM架構以遠低于硬件成本提供了相似的訓練性能(性能差距在7%以內),這將在第6.4節中進行評估。
6.3 機架間架構探索
圖18:機架間網絡架構
在機架內2D-FM之外,UB-Mesh在UB-Mesh Pod內采用另一個2D-FM,形成4D-FM網絡架構。如圖18所示,我們比較了UB-Mesh的2D全互連機架間架構與基線Clos架構:
(a) 2D-FM:16個機架在水平和垂直方向直接互連,形成2D全互連直接連接。所有機架連接到HRS交換機,支持跨Pod互連。在UB-Mesh-Pod中,可支持三種路由策略:
- 最短路徑(Shortest):基線路由策略,僅選擇2D網格上的最短路徑用于P2P通信。
- 繞行(Detour):根據全路徑路由機制,還使用其他路徑以最大化帶寬。
- 借用(Borrow):由于所有機架也連接到交換機,我們允許機架“借用”交換機的帶寬。
(b) Clos:所有機架連接到HRS交換機,無直接連接。此架構比2D-FM消耗更多交換機,但提供最高的All-to-All帶寬和靈活性。
圖19:不同機架間互連的端到端性能比較
性能比較:如圖19所示,我們比較了2D-FM(包括不同路由策略)和Clos架構的性能。可以看出,與理想的Clos架構相比,2D-FM與Clos之間的性能差距微乎其微,尤其是應用Detour和Borrow策略時。GPT3-175B對機架間通信性能不敏感,而在GPT4-2T上,采用最短路徑路由的樸素2D-FM訓練性能下降0.73%。采用Detour和Borrow路由策略后,性能差距縮小至僅0.46%。總體而言,2D-FM機架間互連表現出與昂貴的Clos架構幾乎相同的性能。
圖20:機架間帶寬探索
機架間帶寬探索:圖20-(b)比較了8000 SuperPod在不同機架間帶寬條件下的吞吐性能。評估的機架間帶寬為每NPU的x4、x8、x16和x32 UB IO。對于8000至3.2萬的序列長度,最優機架間帶寬為UB x16;而對于6.4萬至1000萬的序列長度,最優機架間帶寬為UB x32。從UB x8增加到UB x16在8000-3.2萬序列長度范圍內的性能增益僅為0.44%,但從UB x16增加到UB x32在6.4萬-1000萬序列長度范圍內的性能增益更為顯著,達1.85%。在6.4萬至1000萬序列長度的場景中,部分TP(張量并行)和SP(序列并行)流量不可避免地穿越機架間鏈路。更高的機架間帶寬顯著減少了TP和SP的通信時間,在這些場景中帶來更明顯的性能提升。這些數據強調了將機架間帶寬與不同模型場景的特定序列長度需求相匹配的重要性,特別是在減少大規模模型的TP和SP通信延遲方面。
UB-Mesh默認分配每NPU UB x16 IO用于機架間通信,以在成本和性能之間實現平衡。我們還可調整機架內/機架間帶寬比例,以匹配特定LLM訓練工作負載的需求。
6.4 成本效益比較
圖21:CapEx比較
系統成本通常以總擁有成本(TCO)衡量,即TCO = CapEx + OpEx。考慮到相對于基線的訓練性能,我們定義系統的成本效益如下:
我們使用內部數據估算CapEx成本,包括NPU、CPU、LRS、HRS、電纜和其他模塊的成本,并比較不同架構。如圖21所示,UB-Mesh的4D-FM+Clos架構相較于2D-FM+x16(表示每NPU UB x16 IO)Clos、1D-FM+x16 Clos和x64T Clos架構分別實現了1.18倍、1.26倍、1.65倍和2.46倍的CapEx降低。與基線Clos架構相比,UB-Mesh成功將系統中網絡基礎設施成本占比從67%降低至20%,這得益于節省了高性能交換機和長距離光纜/模塊。根據我們的評估,與基線Clos架構相比,98%的高基數交換機和93%的光模塊得以節省。
OpEx降低主要包括系統生命周期內的電費和維護成本。UB-Mesh相較于Clos減少了約35%的OpEx,這得益于其大幅減少了交換機和光模塊的使用。根據我們云部門的AI系統估算,OpEx約占TCO的30%。最終,根據公式(1),UB-Mesh實現了2.04倍更高的成本效益。
6.5 線性度評估
圖22:序列長度256K時的線性度分析
AI集群的線性度指的是集群性能隨NPU數量增加的線性擴展程度。具體而言,我們通過以下公式測量線性度:
圖22評估了UB-Mesh在不同集群規模下的線性度。基準規模(即圖中的1×規模)因任務而異。具體而言,Llama2-70B使用128個NPU,GPT3-175B的基準規模為512,Dense-1T和GPT4-2T使用1000個NPU。
可以看出,UB-Mesh在所有任務的1×至32×規模下線性度超過100%,這是因為規模提升提供了更多高帶寬域,并釋放了搜索更優并行策略以提升MFU(模型浮點利用率)的潛力。GPT4-2T和Dense-1T模型在64×規模(涉及6.4萬個NPU)下線性度有所下降,但仍高于95%。
6.6 網絡可靠性分析
AI訓練過程對硬件故障敏感。經典的基于時間的指標用于評估系統可靠性:
其中,MTBF表示平均故障間隔時間,MTTR表示平均修復時間。
表6:MTBF估算
如表6估算,優先使用直接連接的電線(E-Cables)而非光纖和交換機,大幅降低了網絡模塊的年化故障率(AFR)。
然后,我們計算兩種架構的MTBF
:基線Clos在8000個NPU集群中為13.8小時,而UB-Mesh達到98.5小時,提升了7.14倍。最終,根據公式(3)(我們根據現有統計數據假設MTTR為75分鐘),UB-Mesh的可用性為98.8%,顯著優于Clos的91.6%(提升7.2%)。
為進一步提升可用性,我們還精心設計了內部網絡監控工具,可在10分鐘內快速識別和定位網絡故障,并在3分鐘內觸發任務遷移,大幅減少MTTR。通過這種MTTR優化,根據我們的評估,UB-Mesh的可用性可進一步提升至99.78%。
總結:與基線Clos架構相比,UB-Mesh以微小的性能下降(7%以內)為代價實現了2.04倍更高的成本效益。由于大幅減少了交換機和光模塊的使用,UB-Mesh將網絡可用性提升了7.2%。UB-Mesh還在多個LLM訓練任務上實現了95%以上的線性度。
7.討論
集合通信協處理器
值得一提的是,在UB IO控制器中還配備了一個特殊的協處理器,稱為集合通信單元(Collective Communication Unit, CCU),用于卸載集合通信任務。具體而言,CCU執行指令,主動從HBM讀取或寫入數據,啟動NPU間傳輸,并利用片上SRAM緩沖區執行在線數據歸約操作。這種設計消除了從應用程序內存緩沖區到通信緩沖區的冗余數據拷貝,有效緩解了HBM帶寬的消耗,并通過基于檢查位的細粒度同步機制維持確定的歸約順序。CCU還能與計算核心無縫協作,實現高效的計算與通信重疊。
支持大規模專家模型
除了密集模型和常規MoE模型外,我們注意到人們正在積極探索包含大規模專家的MoE模型[6, 7]。這種設計通常需要大規模、細粒度的All-to-All通信。UB-Mesh通過多路徑和層次化的All-to-All優化(見第5.1節)以及基于UB的Load/Store數據傳輸,能夠高效支持此類大規模專家模型。CCU還能有效卸載All-to-All操作,節省寶貴的計算核心使用,這與[7]中的預期一致。
8.結論
本文提出了UB-Mesh,一種為下一代大規模語言模型(LLM)訓練設計的新型數據中心網絡架構。UB-Mesh采用n維全互連(nD-FullMesh)網絡拓撲,減少了交換機和光模塊的使用,同時適配LLM工作負載的分層局部化流量模式。我們深入考慮了架構設計、物理實現和網絡系統優化,并提出了多項技術以應對各種挑戰。與傳統的Clos網絡相比,UB-Mesh在提供相似LLM訓練性能的同時,實現了2.04倍的系統級成本效益提升和7.2%的可用性提升。此外,UB-Mesh在多項LLM訓練任務上實現了95%以上的線性擴展能力。
本文轉載自??Andy730??,作者:常華Andy
