• ?標(biāo)題：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

摘要

隨著大規(guī)模語言模型（LLM）的規(guī)模持續(xù)擴(kuò)大，其所需的計(jì)算能力和帶寬需求也在不斷提升。為應(yīng)對這一挑戰(zhàn)，我們提出了UB-Mesh，一種新穎的AI數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)，旨在提升可擴(kuò)展性、性能、成本效益以及可用性。與傳統(tǒng)數(shù)據(jù)中心提供對稱的節(jié)點(diǎn)間帶寬不同，UB-Mesh采用了一種分層局部化的n維全互連（nD-FullMesh）網(wǎng)絡(luò)拓?fù)浣Y(jié)構(gòu)。該設(shè)計(jì)充分利用了LLM訓(xùn)練中的數(shù)據(jù)局部性，優(yōu)先采用短距離直接互連，以最小化數(shù)據(jù)移動距離并減少交換機(jī)的使用。

盡管UB-Mesh的n維全互連拓?fù)湓诶碚撋暇哂卸囗?xiàng)優(yōu)勢，但其具體的架構(gòu)設(shè)計(jì)、物理實(shí)現(xiàn)以及網(wǎng)絡(luò)系統(tǒng)優(yōu)化帶來了新的挑戰(zhàn)。為實(shí)現(xiàn)UB-Mesh的實(shí)際構(gòu)建，我們首先設(shè)計(jì)了基于4維全互連拓?fù)涞腢B-Mesh-Pod架構(gòu)。UB-Mesh-Pod通過一系列硬件組件實(shí)現(xiàn)，這些組件構(gòu)成了基礎(chǔ)構(gòu)建模塊，包括專門設(shè)計(jì)的NPU、CPU、低基數(shù)交換機(jī)（LRS）、高基數(shù)交換機(jī)（HRS）、網(wǎng)卡（NIC）等。這些組件通過一種新穎的統(tǒng)一總線（Unified Bus, UB）技術(shù)實(shí)現(xiàn)互連，該技術(shù)支持靈活的IO帶寬分配和硬件資源池化。在網(wǎng)絡(luò)系統(tǒng)優(yōu)化方面，我們提出了名為全路徑路由（All-Path-Routing, APR）的先進(jìn)路由機(jī)制，以高效管理數(shù)據(jù)流量。這些優(yōu)化措施結(jié)合拓?fù)涓兄男阅茉鰪?qiáng)和諸如64+1備份設(shè)計(jì)的穩(wěn)健可靠性機(jī)制，使得UB-Mesh相比傳統(tǒng)的Clos架構(gòu)實(shí)現(xiàn)了2.04倍的成本效益提升、7.2%的網(wǎng)絡(luò)可用性提升，并在多種LLM訓(xùn)練任務(wù)中實(shí)現(xiàn)了95%以上的線性擴(kuò)展能力。。

1.引言

新興的大規(guī)模語言模型（LLM）[1, 4, 6, 8, 10, 18, 22, 26, 27, 31, 32]正在改變AI產(chǎn)業(yè)和人類社會。根據(jù)規(guī)模定律（Scaling Laws）[5, 13]，LLM通過增加模型參數(shù)和訓(xùn)練數(shù)據(jù)量持續(xù)提升其理解、生成和推理能力。然而，這一趨勢對底層訓(xùn)練系統(tǒng)和基礎(chǔ)設(shè)施提出了日益嚴(yán)峻的挑戰(zhàn)，迫使下一代AI數(shù)據(jù)中心滿足以下要求：

圖1：傳統(tǒng)Clos數(shù)據(jù)中心架構(gòu)與UB-Mesh的比較

R1：大規(guī)模。隨著模型規(guī)模和訓(xùn)練數(shù)據(jù)量的增加，需要越來越多的NPU（神經(jīng)網(wǎng)絡(luò)處理單元）或GPU在合理的時間內(nèi)完成訓(xùn)練。例如，Llama-3的預(yù)訓(xùn)練使用1.6萬個GPU需要54天[8]。領(lǐng)先公司最近的公告顯示，已成功部署了擁有10萬GPU的AI訓(xùn)練系統(tǒng)[30]?？蓴U(kuò)展的基礎(chǔ)設(shè)施對于支持LLM技術(shù)的持續(xù)發(fā)展至關(guān)重要。

R2：高帶寬。在LLM訓(xùn)練系統(tǒng)中，AI計(jì)算節(jié)點(diǎn)（NPU/GPU）需要的互連帶寬超過每節(jié)點(diǎn)3.2 Tbps[12, 17]，大約是當(dāng)代數(shù)據(jù)中心中典型CPU節(jié)點(diǎn)互連帶寬的10倍。因此，最先進(jìn)的AI訓(xùn)練系統(tǒng)的總帶寬比當(dāng)前基于CPU的基礎(chǔ)設(shè)施即服務(wù)（IaaS）系統(tǒng)高出10到100倍。

R3：成本效益。構(gòu)建大規(guī)模AI數(shù)據(jù)中心需要巨額硬件投資，資本支出（CapEx）往往達(dá)到數(shù)十億美元。為實(shí)現(xiàn)所需的10到100倍總互連帶寬增長，若使用傳統(tǒng)的對稱Clos數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)（見圖1-(a)），互連成本也將增加10到100倍。優(yōu)化網(wǎng)絡(luò)基礎(chǔ)設(shè)施為提升成本效益提供了重要機(jī)會。此外，降低運(yùn)營成本（OpEx），包括能耗和維護(hù)費(fèi)用，對于確保整體成本效益同樣至關(guān)重要。

R4：高可用性。擁有10萬計(jì)算節(jié)點(diǎn)和約100萬個光模塊的大型LLM訓(xùn)練集群面臨顯著的可用性挑戰(zhàn)。當(dāng)前統(tǒng)計(jì)數(shù)據(jù)顯示，即使每條鏈路的平均故障間隔時間（MTBF）為5年，整個10萬GPU的AI集群的原始MTBF仍下降至不到30分鐘[24]。為此，網(wǎng)絡(luò)架構(gòu)設(shè)計(jì)不僅需提升硬件可靠性，還需融入容錯機(jī)制，以應(yīng)對互連、計(jì)算資源、控制系統(tǒng)和存儲的故障。

同時實(shí)現(xiàn)這些目標(biāo)極具挑戰(zhàn)性，需要向先進(jìn)的下一代數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)設(shè)計(jì)轉(zhuǎn)變。我們認(rèn)為，設(shè)計(jì)下一代AI數(shù)據(jù)中心應(yīng)以以下三個原則為核心：

P1：流量模式驅(qū)動的網(wǎng)絡(luò)拓?fù)?。與傳統(tǒng)數(shù)據(jù)中心工作負(fù)載通常產(chǎn)生均勻且隨機(jī)的流量不同，大規(guī)模語言模型（LLM）的訓(xùn)練流量具有確定性并表現(xiàn)出強(qiáng)烈的數(shù)據(jù)局部性。例如，張量并行（Tensor Parallelism）[25]操作所需的集合通信通常占總數(shù)據(jù)流量的50%以上，且通常發(fā)生在8-64個相鄰NPU的集群內(nèi)。相比之下，數(shù)據(jù)并行（Data Parallelism）產(chǎn)生的集合通信占總流量的不到2%，但通常需要長距離傳輸。因此，層次化、局部化的網(wǎng)絡(luò)架構(gòu)對于適配這些流量模式至關(guān)重要。

P2：拓?fù)涓兄挠?jì)算與通信。在分層局部化的數(shù)據(jù)中心網(wǎng)絡(luò)中，有效運(yùn)行LLM訓(xùn)練是另一重大挑戰(zhàn)。如果訓(xùn)練任務(wù)未在計(jì)算資源上優(yōu)化分布，或網(wǎng)絡(luò)系統(tǒng)未得到充分優(yōu)化，AI集群可能因流量擁塞或帶寬利用不足而性能低下。為解決這一問題，并行策略選擇、路由、集合通信、負(fù)載均衡等因素必須與網(wǎng)絡(luò)拓?fù)渚_對齊。

P3：容錯自愈系統(tǒng)。LLM訓(xùn)練系統(tǒng)必須具備自愈能力以確保魯棒性。鏈路故障時，路由系統(tǒng)應(yīng)自動切換至備用路徑；同樣，若NPU發(fā)生故障，應(yīng)有機(jī)制無縫激活備用NPU，以維持系統(tǒng)完整性和LLM訓(xùn)練的連續(xù)性。

為滿足要求R1-R4并遵循原則P1-P3，我們提出了創(chuàng)新的UB-Mesh架構(gòu)。如圖1-(b)所示，UB-Mesh采用n維全互連（nD-FullMesh）網(wǎng)絡(luò)拓?fù)?，通過遞歸構(gòu)建全互連拓?fù)洹獜陌迳舷噜廚PU間的1維連接開始，擴(kuò)展到機(jī)架內(nèi)相鄰1維網(wǎng)格間的2維連接，再通過互連更廣范圍的相鄰高維網(wǎng)格擴(kuò)展至3維及以上。這種分層局部化的網(wǎng)絡(luò)架構(gòu)最小化傳輸跳數(shù)并優(yōu)化每跳距離，優(yōu)先采用直接互連而非間接長距離交換，從而減少對交換機(jī)和光模塊的依賴，滿足R1和R3。此外，它根據(jù)LLM訓(xùn)練的傳輸需求實(shí)現(xiàn)層次化帶寬分配，為短距離通信提供高帶寬，為長距離通信提供較低帶寬，滿足R2和R3并遵循P1。

遵循P2，我們深入研究了先進(jìn)的網(wǎng)絡(luò)和系統(tǒng)優(yōu)化機(jī)制以增強(qiáng)UB-Mesh架構(gòu)。具體而言，我們引入了全路徑路由（All-Path Routing, APR）技術(shù)，以充分利用直接連接鏈路的帶寬。APR結(jié)合源路由、結(jié)構(gòu)化尋址與線性表查找以及無死鎖流量控制機(jī)制，實(shí)現(xiàn)自適應(yīng)路由、最小化轉(zhuǎn)發(fā)開銷并避免死鎖。此外，我們加入了拓?fù)涓兄目焖俟收匣謴?fù)機(jī)制以提升可靠性。為進(jìn)一步優(yōu)化性能，我們提出了拓?fù)涓兄募贤ㄐ藕筒⑿袃?yōu)化，以提升訓(xùn)練期間的帶寬利用率。

為滿足R4并遵循P3，UB-Mesh采用64+1高可用性設(shè)計(jì)：每個機(jī)架包含一個額外的備用NPU。當(dāng)系統(tǒng)中NPU發(fā)生意外故障時，備用NPU被激活以恢復(fù)功能，確保LLM訓(xùn)練任務(wù)不間斷繼續(xù)。此外，路由系統(tǒng)通過一種新穎的直接通知技術(shù)在鏈路故障時實(shí)現(xiàn)快速恢復(fù)。

我們精心設(shè)計(jì)了UB-Mesh的硬件和系統(tǒng)堆棧，綜合考慮了各種工程約束和權(quán)衡。具體實(shí)現(xiàn)中，UB-Mesh-Pod采用4維全互連拓?fù)?，使UB-Mesh可無縫擴(kuò)展至8000個NPU，形成支持下一代AI數(shù)據(jù)中心構(gòu)建的高帶寬域。為此，我們開發(fā)了一系列硬件組件作為基礎(chǔ)構(gòu)建模塊，包括NPU、CPU、低基數(shù)交換機(jī)（LRS）、高基數(shù)交換機(jī)（HRS）、網(wǎng)卡（NIC）等。

與采用多種互連技術(shù)（如PCIe、NVLINK、IB和RoCE）的基線系統(tǒng)不同，UB-Mesh使用新穎的統(tǒng)一總線（Unified Bus, UB）技術(shù)實(shí)現(xiàn)所有組件互連。這種統(tǒng)一方法提升了IO資源分配的靈活性，其點(diǎn)對點(diǎn)通信能力支持高效的硬件資源池化，UB還為無縫跨層優(yōu)化提供了機(jī)會。

全面評估表明，與非超訂閱的Clos網(wǎng)絡(luò)相比，UB-Mesh將高基數(shù)交換機(jī)使用量減少98%，光模塊使用量減少93%，實(shí)現(xiàn)系統(tǒng)級成本效益提升2.04倍。在多個LLM訓(xùn)練任務(wù)中的實(shí)驗(yàn)還顯示，與成本高昂的Clos網(wǎng)絡(luò)相比，UB-Mesh的性能降幅僅在7%以內(nèi)。這種低成本與高性能的結(jié)合不僅滿足當(dāng)前LLM訓(xùn)練需求，還使其能夠有效應(yīng)對未來的可擴(kuò)展性挑戰(zhàn)。

2.背景與動機(jī)

2.1 LLM訓(xùn)練中的“通信墻”

大規(guī)模語言模型（LLM）訓(xùn)練是有史以來規(guī)模最大、計(jì)算和通信需求最為密集的并行計(jì)算應(yīng)用[11, 18, 20–22, 25, 27, 32]。遵循所謂的規(guī)模定律（Scaling Laws）[3, 5, 13]，LLM通過增加模型參數(shù)和訓(xùn)練數(shù)據(jù)量來提升性能。因此，LLM需要越來越多的AI加速器以在合理的時間內(nèi)完成訓(xùn)練。例如，開源的Llama-3.1模型在1.6萬個GPU上進(jìn)行訓(xùn)練[8]，而下一代LLM模型已使用10萬GPU進(jìn)行訓(xùn)練[24, 29]。

標(biāo)準(zhǔn)的訓(xùn)練過程包括重復(fù)的訓(xùn)練迭代，每個迭代包含前向傳播以計(jì)算損失、反向傳播以確定梯度，以及優(yōu)化器步驟以調(diào)整模型參數(shù)。為充分利用分布式計(jì)算能力，LLM訓(xùn)練通過多種并行策略將數(shù)據(jù)、模型和激活值分割到數(shù)萬個NPU上。在每個迭代中，NPU之間頻繁交換數(shù)據(jù)以分發(fā)輸入數(shù)據(jù)、同步激活值和梯度等。隨著訓(xùn)練系統(tǒng)規(guī)模的擴(kuò)大，數(shù)據(jù)移動成為系統(tǒng)中成本最高的部分[9, 19]。若缺乏強(qiáng)大的NPU間通信能力，訓(xùn)練過程很容易受限于“通信墻”（Communication Wall）。

2.2 LLM訓(xùn)練中數(shù)據(jù)流量的局部性

圖2：LLM訓(xùn)練中的并行性

如圖2所示，LLM訓(xùn)練通常涉及多種并行技術(shù)，以下逐一介紹：

張量并行（Tensor Parallelism, TP）：TP以行或列的方式分割模型層，并將子層放置在多個NPU上并行計(jì)算[25]。其主要涉及AllReduce操作以合并分布式的部分結(jié)果。

序列并行（Sequence Parallelism, SP）：SP（在某些論文中也稱為上下文并行）通常用于將序列分割到多個NPU以實(shí)現(xiàn)并行處理。SP依賴RingAttention[15]技術(shù)，并可采用AllGather操作收集不同NPU的部分結(jié)果。

專家并行（Expert Parallelism, EP）：對于采用專家混合（Mix-of-Experts, MoE）技術(shù)的LLM模型，密集MLP層被MoE層替代，每個MoE層包含若干“專家”。執(zhí)行時專家被稀疏激活。EP將專家分布在不同NPU上，輸入標(biāo)記通過All2All通信動態(tài)發(fā)送至目標(biāo)專家。

流水線并行（Pipeline Parallelism, PP）：與TP分割每個模型層不同，PP將層分布到多個設(shè)備上，并以流水線方式執(zhí)行前向和反向傳播。PP涉及低開銷的P2P通信以傳輸層間激活值，但需要高效的調(diào)度算法[11, 16]以最小化流水線中的空泡。

數(shù)據(jù)并行（Data Parallelism, DP）：DP在多個NPU上復(fù)制模型和優(yōu)化器狀態(tài)，每個副本并行處理一部分輸入批次。訓(xùn)練過程中需通過AllReduce操作同步梯度。

表1：LLM訓(xùn)練中數(shù)據(jù)流量的分析

這些并行技術(shù)共同將訓(xùn)練任務(wù)分布到數(shù)千個NPU上。需要注意的是，并非所有并行技術(shù)產(chǎn)生的數(shù)據(jù)流量均等。根據(jù)我們基于內(nèi)部MoE-2T模型的分析（如表1詳述），通信強(qiáng)度呈層次化并表現(xiàn)出強(qiáng)烈的局部性。具體而言，TP和SP約占總流量的97%，而其余并行技術(shù)通常產(chǎn)生不到2%的流量。其他模型架構(gòu)可能表現(xiàn)出略有不同的數(shù)據(jù)流量分布，但同樣具有較強(qiáng)的局部性[28]。因此，架構(gòu)設(shè)計(jì)應(yīng)優(yōu)先采用層次化的網(wǎng)絡(luò)帶寬分配方式。

2.3 數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)

圖3：傳統(tǒng)數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)

設(shè)計(jì)大規(guī)模LLM訓(xùn)練集群涉及一個關(guān)鍵問題：如何將眾多NPU資源組織成一個協(xié)調(diào)的分布式系統(tǒng)？這需要確定合適的數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)。傳統(tǒng)數(shù)據(jù)中心和超算系統(tǒng)探索了多種拓?fù)浣Y(jié)構(gòu)，但這些拓?fù)湮幢剡m用于大規(guī)模LLM訓(xùn)練。圖3展示了一些常見拓?fù)洌?/p>

• Clos：傳統(tǒng)數(shù)據(jù)中心通常采用Clos架構(gòu)，通過兩到三層交換機(jī)對稱連接NPU/CPU，提供高性能并適應(yīng)多種流量模式。然而，其高成本是一個缺點(diǎn)，主要源于對高性能交換機(jī)和光模塊的廣泛使用。
• 3D Torus：與依賴交換機(jī)互連的Clos不同，3D Torus直接連接相鄰NPU。雖然降低了成本，但3D Torus的NPU間帶寬較低，且難以適應(yīng)復(fù)雜的通信模式，例如MoE模型中常見的All-to-All通信[20]。
• Dragon-Fly：Dragon-Fly拓?fù)鋄14]將交換機(jī)組織為互連組，每組內(nèi)交換機(jī)彼此直接相連，組間通過長距離電纜連接，數(shù)據(jù)包通過間接路由傳輸。該架構(gòu)在傳統(tǒng)數(shù)據(jù)中心和高性能計(jì)算（HPC）環(huán)境中降低了維度。Dragon-Fly比Clos成本更低，但因需滿足完整的NPU-交換機(jī)帶寬需求仍較昂貴，且其架構(gòu)在LLM訓(xùn)練流量（特別是P2P和AllReduce場景）中表現(xiàn)不佳。
• Fugaku Tofu：Fugaku Tofu[2]提出了一種獨(dú)特的6D Torus拓?fù)溆糜贖PC應(yīng)用。然而，與3D Torus類似，其NPU間帶寬較低，且難以適應(yīng)復(fù)雜的通信模式。
• 綜上所述，如何為大規(guī)模LLM訓(xùn)練設(shè)計(jì)一個高性能且成本效益高的數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)仍是一個開放性問題。

3.UB-Mesh架構(gòu)

為滿足第1節(jié)中提出的要求R1-R4并遵循原則P1-P3，我們提出了UB-Mesh架構(gòu)。UB-Mesh采用了一種新穎的n維全互連（nD-FullMesh）拓?fù)浣Y(jié)構(gòu)，最大限度地利用直接電線互連，從而減少對昂貴高帶寬交換機(jī)和光模塊的需求。接下來，我們將介紹一系列專為構(gòu)建UB-Mesh設(shè)計(jì)的硬件模塊。最后，我們將詳細(xì)討論UB-Mesh-Pod的架構(gòu)設(shè)計(jì)，即UB-Mesh的具體4維全互連實(shí)現(xiàn)，同時充分考慮了若干工程約束和權(quán)衡。

3.1 nD-FullMesh拓?fù)?/h4>

圖4：UB-Mesh的n維全互連拓?fù)浼翱赡艿奈锢碛成?/p>

如圖4所示，UB-Mesh引入了一種新穎的n維全互連拓?fù)洌瑥?qiáng)調(diào)跨所有網(wǎng)絡(luò)層級的直接互連。該拓?fù)鋸南噜徆?jié)點(diǎn)形成1維全互連開始，其中同一層級的每個節(jié)點(diǎn)與該層內(nèi)其他所有節(jié)點(diǎn)相連。基于這一概念，相鄰1維全互連的節(jié)點(diǎn)之間再進(jìn)行直接互連，形成2維全互連。這一過程遞歸進(jìn)行，相鄰的2維全互連形成3維全互連，依此類推，最終根據(jù)需要擴(kuò)展至n維全互連。這種靈活的虛擬拓?fù)淇梢詿o縫映射到各種物理NPU組織形式：單板上的1維全互連、機(jī)架內(nèi)的2維全互連、機(jī)架行間的3維全互連、樓層機(jī)架組內(nèi)的4維全互連，甚至跨越整棟建筑的5維全互連及更高維度。

圖5：高維拓?fù)渑c靈活的帶寬分配

這種分層局部化的拓?fù)湓跇?gòu)建下一代LLM訓(xùn)練數(shù)據(jù)中心時具有多項(xiàng)優(yōu)勢。首先，由于nD-FullMesh在每個網(wǎng)絡(luò)維度上形成緊密耦合的直接連接域，并可提供逐層超訂閱的帶寬，它充分利用了LLM訓(xùn)練固有的數(shù)據(jù)局部性和從密集到稀疏的流量模式（見第2.2節(jié)）。此外，我們可以靈活調(diào)整每個維度的每節(jié)點(diǎn)帶寬分配，以滿足未來LLM訓(xùn)練任務(wù)的具體需求。例如，如圖5-(a)所示，假設(shè)存在6維全互連拓?fù)?，我們可以通過在NPU的IO模塊中分配不同數(shù)量的互連鏈路來調(diào)整??????和??????維度的帶寬，如圖5-(b)所示。

其次，與需要大量交換的Clos架構(gòu)不同，這種直接連接的拓?fù)湓跍p少傳輸距離方面具有巨大潛力。結(jié)合數(shù)據(jù)放置和集合通信優(yōu)化（見第5.2節(jié)），大多數(shù)傳輸（通常涉及TP或SP，如表1所示）可在0-2跳內(nèi)完成，大幅降低數(shù)據(jù)移動開銷。

表2：不同類型鏈路的使用估算

第三，與Clos等其他流行拓?fù)湎啾?，nD-FullMesh拓?fù)錅p輕了對高帶寬交換機(jī)和長距離光互連的依賴。相反，它最大限度地利用短距離直接互連。例如，根據(jù)表2的估算，短距離無源電線占據(jù)了總電纜消耗的86.7%。這一設(shè)計(jì)顯著降低了交換機(jī)和光模塊的成本，同時提升了系統(tǒng)可靠性，因?yàn)殡娋€和連接器比光模塊更穩(wěn)定。

最后，與不適合復(fù)雜集合通信操作（如All-to-All）的3D-Torus拓?fù)洳煌?，全互連拓?fù)淠軌蚋咝еС诌@些操作，具體將在第5節(jié)中討論。

3.2 UB-Mesh的構(gòu)建模塊

所提出的nD-FullMesh拓?fù)湓诶碚撋暇哂卸囗?xiàng)優(yōu)勢，但其實(shí)踐實(shí)現(xiàn)和系統(tǒng)性能的提升帶來了新的挑戰(zhàn)，需要全面考慮各種工程約束和權(quán)衡。為此，我們設(shè)計(jì)并制造了一系列硬件模塊，作為UB-Mesh的基礎(chǔ)構(gòu)建塊。這些模塊通過一種名為統(tǒng)一總線（Unified Bus, UB）的新型互連機(jī)制實(shí)現(xiàn)連接。具體細(xì)節(jié)如下：

3.2.1 硬件模塊

表3：UB-Mesh的主要構(gòu)建模塊

表3列舉了UB-Mesh的主要組件。其核心是作為AI計(jì)算單元的NPU。NPU配備兩個UB IO控制器，總計(jì)提供x72條UB通道。負(fù)責(zé)執(zhí)行主機(jī)程序的CPU也配備一個UB IO控制器，提供x32的UB IO。

盡管UB-Mesh優(yōu)先采用直接互連，但交換能力對于滿足特定需求仍不可或缺。例如，在圖5中聚合??和??????維度的NPU IO有助于實(shí)現(xiàn)高效的機(jī)架間互連。此外，我們期望CPU-NPU比例及其綁定關(guān)系能夠靈活調(diào)整，以實(shí)現(xiàn)NPU和內(nèi)存資源的高效池化，這可以通過交換實(shí)現(xiàn)。

為在UB-Mesh-Pod內(nèi)提供成本效益高的交換能力，我們開發(fā)了輕量級低基數(shù)交換機(jī)（LRS）。其制造成本低，可高效聚合機(jī)架間IO帶寬并支持CPU-NPU通信。

盡管UB-Mesh僅使用低基數(shù)交換機(jī)即可實(shí)現(xiàn)，但實(shí)際網(wǎng)絡(luò)約束通常將其規(guī)模限制在一定范圍內(nèi)（UB-Mesh-Pod）。因此，UB-Mesh還包括高基數(shù)交換機(jī)（HRS），提供x512的UB IO容量以支持Pod級交換。UB-Mesh的詳細(xì)架構(gòu)設(shè)計(jì)和考慮將在第3.3節(jié)中介紹。

3.2.2 統(tǒng)一總線（UB）互連

圖6：統(tǒng)一總線互連與傳統(tǒng)混合互連的對比

連接這些異構(gòu)硬件模塊需要特定的互連技術(shù)。如圖6-(a)所示，傳統(tǒng)的基于GPU的LLM系統(tǒng)通常采用多種互連方式，而UB-Mesh通過統(tǒng)一總線（UB）技術(shù)實(shí)現(xiàn)所有組件的互連，其優(yōu)勢包括：  

• 靈活的IO資源分配：UB互連與特定用例解耦，支持芯片內(nèi)不同類型IO的靈活資源分配，如圖5所示。NPU間帶寬和CPU-NPU帶寬也可根據(jù)具體需求靈活調(diào)整，因?yàn)樗鼈兪褂孟嗤腢B鏈路。
• 硬件資源池化：UB的點(diǎn)對點(diǎn)通信能力支持包括DDR DRAM、CPU、NPU和NIC在內(nèi)的高效硬件資源池化。例如，在圖6-(b)中，CPU和NPU通過UB互連實(shí)現(xiàn)池化，以提升資源利用率。
• 系統(tǒng)優(yōu)化收益：統(tǒng)一互連消除了協(xié)議轉(zhuǎn)換的需要，大幅降低開銷并簡化驅(qū)動程序、通信庫和框架的設(shè)計(jì)與優(yōu)化。

3.3 架構(gòu)概覽

圖7：UB-Mesh的架構(gòu)設(shè)計(jì)

通過組合表3中列出的硬件構(gòu)建模塊并通過統(tǒng)一總線（Unified Bus, UB）進(jìn)行互連，我們提出了UB-Mesh的整體架構(gòu)設(shè)計(jì)。如圖7-(a)所示，我們按照所提出的nD-FullMesh拓?fù)鋵?shí)現(xiàn)了一個UB-Mesh-Pod。具體而言，在該架構(gòu)中，我們在每個機(jī)架內(nèi)創(chuàng)建了一個2D全互連拓?fù)?，并將其擴(kuò)展到機(jī)架之外的另一個2D全互連拓?fù)洌瑥亩纬梢粋€4D全互連拓?fù)?。我們選擇不在這一代產(chǎn)品中將UB-Mesh擴(kuò)展至5D全互連或更高維度，以在成本效益和靈活性之間實(shí)現(xiàn)工程平衡，在4D全互連之外仍采用Clos拓?fù)?。以下為詳?xì)介紹：

3.3.1 在機(jī)架內(nèi)實(shí)現(xiàn)2D全互連

圖8：UB-Mesh機(jī)架的硬件實(shí)現(xiàn)

如圖7-(b)所示，每個機(jī)架配置為2D全互連拓?fù)?，集成了多個組件以確保高效通信和資源利用。機(jī)架的核心由八個NPU板組成，每個板包含八個NPU，如圖8-(a)所示。這64個NPU在機(jī)架內(nèi)互連，形成一個2D全互連網(wǎng)絡(luò)，確保高NPU間帶寬。需要注意的是，由于UB IO控制器也具備路由能力，每個NPU在此架構(gòu)中既充當(dāng)計(jì)算單元又充當(dāng)路由器，支持間接路由。

除NPU外，機(jī)架還包括專用CPU板，如圖8-(c)所示。與傳統(tǒng)設(shè)置中CPU和NPU位于同一板不同，此處它們是分離的。CPU通過交換機(jī)連接到NPU，支持靈活的CPU/NPU比例調(diào)整，并實(shí)現(xiàn)CPU/NPU/DDR資源池化以提升資源利用率。

機(jī)架配備多個背板交換機(jī)，管理機(jī)架內(nèi)和機(jī)架間的連接。這些交換機(jī)采用低基數(shù)設(shè)計(jì)（稱為LRS），以降低成本，同時確保設(shè)備間的無阻塞通信?？傮w而言，這些背板交換機(jī)輸出四個UB x256 IO。

3.3.2 64+1高可用性設(shè)計(jì)

正如第1節(jié)所述，在大規(guī)模AI數(shù)據(jù)中心中確保系統(tǒng)可用性是一個關(guān)鍵挑戰(zhàn)。系統(tǒng)中單個NPU可能會發(fā)生意外故障。雖然軟件系統(tǒng)可以檢測這些故障并使用剩余健康NPU重啟訓(xùn)練，但由于計(jì)算能力和系統(tǒng)帶寬的減少，系統(tǒng)性能將受到顯著影響。

圖9：通過啟用備用NPU實(shí)現(xiàn)容錯

為提升系統(tǒng)可用性，我們引入了特殊的64+1設(shè)計(jì)：如圖7-(b)和圖8-(b)所示，系統(tǒng)包含64個常規(guī)NPU以及一個額外的備用NPU。該備用NPU通過LRS與64個常規(guī)NPU相連。當(dāng)任何NPU發(fā)生故障時，備用NPU將被策略性地啟用。如圖9所示，當(dāng)NPU-3發(fā)生故障時，管理系統(tǒng)激活備用NPU（圖中節(jié)點(diǎn)B）以替換NPU-3。與NPU-3相關(guān)的原始直接連接鏈路將被重定向。例如，路徑5-3被重定向?yàn)槁窂?-LRS-B。盡管此策略將原始直接連接變?yōu)閱翁酚?，略微增加了傳輸延遲，但相比簡單屏蔽NPU-3并在剩余七個NPU上運(yùn)行任務(wù)，其效果要優(yōu)越得多。在機(jī)架內(nèi)帶寬為主要關(guān)注點(diǎn)的LLM訓(xùn)練場景中，增加的延遲對整體訓(xùn)練性能影響微乎其微。

3.3.3 在UB-Mesh-Pod中擴(kuò)展至4D全互連

在機(jī)架內(nèi)，我們將64個NPU組織為2D全互連拓?fù)?。在機(jī)架之外，我們實(shí)現(xiàn)另一個2D全互連來組織機(jī)架，形成名為UB-Mesh-Pod的4D全互連架構(gòu)。

如圖7-(a)和(c)所示，一行中四個相鄰機(jī)架被組織為緊密耦合的1D全互連，機(jī)架通過每個機(jī)架中的LRS端口直接互連。在每個機(jī)架列中，機(jī)架也直接相連，形成一個16機(jī)架的2D全互連拓?fù)?。此配置中的每條鏈路代表一個UB x128 IO，如圖8-(d)所示。我們沿兩個維度連接四個相鄰機(jī)架以構(gòu)建機(jī)架間全互連，這是考慮到有源電線的覆蓋范圍得出的最佳點(diǎn)。由于每個機(jī)架有64個NPU，每個Pod有16個機(jī)架，一個4D全互連的UB-Mesh-Pod總共包含1024個NPU。

3.3.4 UB-Mesh-SuperPod及更高規(guī)模

基于1000規(guī)模的UB-Mesh-Pod，我們進(jìn)一步構(gòu)建了UB-Mesh-SuperPod，可容納多個UB-Mesh-Pod?？紤]到當(dāng)前云場景中小規(guī)模或中規(guī)模LLM訓(xùn)練負(fù)載可能不會消耗整個SuperPod，我們選擇在Pod級互連中采用對稱的Clos拓?fù)洌抢^續(xù)使用全互連設(shè)計(jì)。這種設(shè)計(jì)允許云管理者根據(jù)用戶需求靈活劃分SuperPod，并保證每個劃分域內(nèi)的全帶寬。如圖7-(c)所示，我們使用高基數(shù)Pod交換機(jī)（HRS）連接SuperPod中的每個機(jī)架，可擴(kuò)展至8000個NPU。

最后，SuperPod中的機(jī)架還通過UB交換機(jī)（方案(a)）或位于CPU板上的NIC（方案(b)）連接到大規(guī)模數(shù)據(jù)中心網(wǎng)絡(luò)（DCN）。DCN域通常支持大規(guī)模數(shù)據(jù)并行訓(xùn)練。DCN交換機(jī)采用Clos拓?fù)浣M織。DCN域可擴(kuò)展至10萬NPU或更多。

4.實(shí)現(xiàn)架構(gòu)

在UB-Mesh-SuperPod中，4D全互連與Clos混合拓?fù)淦惹行枰咝У幕旌下酚芍С郑撼送ㄟ^LRS和HRS進(jìn)行的基于交換機(jī)的路由外，NPU本身通過UB控制器也具備路由能力。UB-Mesh中大量路由器和層次化拓?fù)錇槁酚上到y(tǒng)引入了新的復(fù)雜性。我們認(rèn)為，路由系統(tǒng)必須滿足以下五個關(guān)鍵要求：  

• 支持混合拓?fù)洌郝酚上到y(tǒng)必須滿足UB-Mesh的4D全互連與Clos混合拓?fù)渌鶐淼男枨蟆?/li>
• 高效轉(zhuǎn)發(fā)：由于每個NPU同時也是路由器，且整個系統(tǒng)包含大量NPU，為節(jié)省NPU硬件資源，路由系統(tǒng)必須高效處理路由表查詢和轉(zhuǎn)發(fā)操作。
• 支持非最短路徑：在nD全互連拓?fù)渲校瑑蓚€端點(diǎn)之間存在多個路徑，距離各異。系統(tǒng)應(yīng)支持使用非最短路徑，以最大化網(wǎng)絡(luò)帶寬利用率。
• 快速故障恢復(fù)：為確?？煽啃院涂捎眯裕酚上到y(tǒng)必須快速從故障中恢復(fù)，減少對訓(xùn)練過程的影響。
• 無死鎖：最后，整個網(wǎng)絡(luò)系統(tǒng)必須在無死鎖風(fēng)險下運(yùn)行，確保數(shù)據(jù)流暢通無阻。

表4：路由系統(tǒng)比較

然而，表4中列出的現(xiàn)有路由技術(shù)，如Clos、Torus和DragonFly拓?fù)渖鲜褂玫拈L前綴匹配（LPM）、基于主機(jī)的路由和維度順序路由（DOR），均無法完全滿足我們的需求。因此，我們提出了全路徑路由（All-Path-Routing, APR）技術(shù)和基于直接通知的快速故障恢復(fù)技術(shù)，以滿足上述要求。

4.1 全路徑路由（APR）

圖10：最短路徑路由與全路徑路由

UB-Mesh的全互連網(wǎng)絡(luò)架構(gòu)為任意兩個NPU之間提供了多條路徑。傳統(tǒng)的路由策略，如最短路徑優(yōu)先（SPF）路由（見圖10-(a)），往往無法充分利用網(wǎng)絡(luò)帶寬，且易受鏈路故障影響。為提升網(wǎng)絡(luò)效率和魯棒性，我們在UB-Mesh中提出了全路徑路由（APR）機(jī)制。

如圖10-(b)所示，APR利用源節(jié)點(diǎn)與目標(biāo)節(jié)點(diǎn)之間的所有可用路徑。這種靈活性允許動態(tài)切換路徑，以應(yīng)對故障或擁塞，從而提升網(wǎng)絡(luò)魯棒性。為高效實(shí)現(xiàn)這一目標(biāo)，APR依賴以下三種底層機(jī)制：(1) 源路由（SR）機(jī)制實(shí)現(xiàn)自適應(yīng)路由；(2) 結(jié)構(gòu)化尋址與線性表查詢技術(shù)，減少路由表查詢和轉(zhuǎn)發(fā)開銷；(3) 拓?fù)涓兄獰o死鎖流量控制，避免死鎖。

4.1.1 源路由

圖11：源路由頭部格式

為充分利用APR提供的路徑，一種實(shí)用方法是采用源路由（SR）機(jī)制。如圖11所示，在原始數(shù)據(jù)包頭部添加一個8字節(jié)的SR頭部，包含轉(zhuǎn)發(fā)指令。每個路由器使用一個4位指針（ptr）指示12位位圖字段中的位偏移。第??位的值指定第??跳的轉(zhuǎn)發(fā)方式（即，位值為1表示該跳采用SR轉(zhuǎn)發(fā)，0表示傳統(tǒng)轉(zhuǎn)發(fā)）。在SR轉(zhuǎn)發(fā)的情況下，位圖字段還用于定位六個指令字段之一，指示如何轉(zhuǎn)發(fā)該數(shù)據(jù)包。SR信息高度壓縮，因此開銷較低。

4.1.2 結(jié)構(gòu)化尋址與線性表查詢

為減少每個NPU的UB IO控制器中的路由表查詢和轉(zhuǎn)發(fā)開銷，APR路由系統(tǒng)利用UB-Mesh拓?fù)涞慕Y(jié)構(gòu)，采用結(jié)構(gòu)化尋址和線性表查詢機(jī)制。具體而言，地址空間根據(jù)網(wǎng)絡(luò)元素的物理位置（如Pod、機(jī)架和板）劃分為多個段。由于同一段內(nèi)的NPU共享相同前綴，僅需存儲短段地址，NPU可通過相對于段地址的線性偏移進(jìn)行尋址。此設(shè)計(jì)顯著減少表空間，加速路由表生成和分發(fā)，提升收斂速度，并支持快速響應(yīng)狀態(tài)變化，例如故障恢復(fù)操作。

4.1.3 拓?fù)涓兄獰o死鎖流量控制

鑒于UB-Mesh的nD全互連拓?fù)浒瑥?fù)雜的環(huán)結(jié)構(gòu)，且APR機(jī)制支持多跳路由，在有限的虛擬通道（VL）資源下實(shí)現(xiàn)無死鎖流量控制具有挑戰(zhàn)性。為應(yīng)對這一挑戰(zhàn)，我們提出了TFC（拓?fù)涓兄獰o死鎖流量控制）算法。該算法在僅使用2個VL資源的情況下，最小化VL資源使用，同時實(shí)現(xiàn)無死鎖的全路徑路由。

TFC算法通過通道依賴圖（CDG）建模死鎖，并將UB-Mesh拓?fù)鋭澐譃樽訄D。在每個子圖內(nèi)，拓?fù)鋵?dǎo)向規(guī)則和VL約束被統(tǒng)一為單一集合。算法應(yīng)用N維跨維度斷環(huán)原則，將集合分解為單VL無環(huán)子集，隨后采用同維度斷環(huán)原則計(jì)算冪集元素的排列和笛卡爾積，生成所有路徑組合和VL映射，確保無死鎖運(yùn)行。由于篇幅限制，詳細(xì)說明在此省略。

4.2 基于直接通知的快速故障恢復(fù)

圖12：從逐跳通知到直接通知

當(dāng)網(wǎng)絡(luò)中發(fā)生鏈路故障時，傳統(tǒng)路由系統(tǒng)通常采用逐跳通知方式。如圖12左側(cè)所示，當(dāng)鏈路1-3發(fā)生故障時，該信息由節(jié)點(diǎn)1和3逐跳傳播，通常耗時較長。在UB-Mesh中，由于每個節(jié)點(diǎn)具有確定的通信目標(biāo)集，我們可以通過直接通知受影響的節(jié)點(diǎn)加速路由收斂。如圖12右側(cè)所示，當(dāng)鏈路1-3發(fā)生故障時，根據(jù)預(yù)計(jì)算的路由關(guān)系，該信息直接發(fā)送至節(jié)點(diǎn)6。通過這種拓?fù)涓兄闹苯油ㄖ绞?，可大幅減少控制平面的開銷。

5.性能最大化

盡管UB-Mesh的硬件架構(gòu)設(shè)計(jì)與LLM訓(xùn)練的流量模式相匹配，但在這種層次化的AI集群上運(yùn)行工作負(fù)載時，如果訓(xùn)練任務(wù)未能在計(jì)算資源上有效分布，可能會導(dǎo)致低利用率。此外，集合通信的性能對整體訓(xùn)練效率至關(guān)重要。為確保LLM訓(xùn)練期間的系統(tǒng)性能最優(yōu)，我們引入了若干拓?fù)涓兄獌?yōu)化策略，以進(jìn)一步提升系統(tǒng)性能。

5.1 拓?fù)涓兄贤ㄐ?/h4>

為優(yōu)化UB-Mesh上的集合通信并充分利用層次化直接互連的帶寬，我們提出了利用UB-Mesh全路徑路由能力（見第4.1節(jié)）的拓?fù)涓兄贤ㄐ潘惴?。我們以All-Reduce和All-to-All為例闡述我們的概念：

圖13：AllReduce的多環(huán)算法

All-Reduce：我們提出了一個多環(huán)（Multi-Ring）算法，以在UB-Mesh上高效實(shí)現(xiàn)All-Reduce。我們首先通過統(tǒng)一的抽象模型對網(wǎng)絡(luò)拓?fù)溥M(jìn)行建模，考慮節(jié)點(diǎn)數(shù)量、節(jié)點(diǎn)間連接、鏈路帶寬和延遲等因素。接下來，我們將集合通信與路徑映射相結(jié)合，采用邏輯多環(huán)算法，確保路徑獨(dú)占使用，避免流量沖突。如圖13-(a)所示，原始路徑表示默認(rèn)映射。未包含在這些路徑中的空閑鏈路通過APR機(jī)制被利用，以提升帶寬。最后，如圖13-(b)所示，我們優(yōu)化了多路徑上的流量分區(qū)，以緩解瓶頸并通過APR最大化借用帶寬的收益。

All-to-All：在All-to-All通信場景中，我們考慮了兩種代表性用例：

圖14：多路徑與層次化All-to-All

(1) 通用All-to-All：如圖14-(a)所示，當(dāng)UB-Mesh中的源節(jié)點(diǎn)（為簡化起見，圖中展示2D網(wǎng)格，提出的技術(shù)可擴(kuò)展至更高維度）向不同目標(biāo)節(jié)點(diǎn)發(fā)送不同數(shù)據(jù)時，我們采用多路徑All2All優(yōu)化。具體而言，每個元素（向量或張量）被分為兩個分區(qū)，分別沿X-FullMesh和Y-FullMesh互連同時傳輸，并最多使用一跳轉(zhuǎn)發(fā)到達(dá)目標(biāo)。這種策略保證了UB-Mesh nD全互連架構(gòu)中的高帶寬利用率。

(2) 廣播+歸約（Broadcast + Reduce）：對于涉及標(biāo)記分發(fā)和專家數(shù)據(jù)收集的All-to-All操作[7]，其語義等價于多個廣播和歸約操作的重疊。如圖14-(b)和(c)所示，我們可采用層次化廣播/歸約以節(jié)省帶寬使用，充分利用UB-Mesh的層次拓?fù)洹?/p>

5.2 拓?fù)涓兄⑿谢?/p>

圖15：拓?fù)涓兄⑿谢?/p>

給定訓(xùn)練任務(wù)和系統(tǒng)，我們確定最優(yōu)并行策略，以充分利用UB-Mesh的高帶寬局部互連。我們采用拓?fù)涓兄⑿谢瘷C(jī)制優(yōu)化LLM訓(xùn)練的模型和數(shù)據(jù)分割。如圖15所示，我們的方法包括：

步驟①：生成可行的并行配置并將其映射到UB-Mesh架構(gòu)上。

步驟②：利用拓?fù)涓兄ㄐ懦杀灸Ｐ驮u估通信成本。

步驟③：迭代最小化通信開銷以找到最優(yōu)配置。

該機(jī)制有兩個重要要求：(1) 構(gòu)建適當(dāng)?shù)乃阉骺臻g以平衡效率和性能；(2) 確保成本模型盡可能準(zhǔn)確。對于要求(1)，我們使用基于優(yōu)先級的啟發(fā)式方法剪枝搜索空間：涉及高通信量的TP和SP（或CP）優(yōu)先分配到高帶寬域，而用于梯度更新的PP和DP優(yōu)先級最低。對于需要EP的MoE模型，我們強(qiáng)制SP*DP為EP的整數(shù)倍。對于要求(2)，我們在UB-Mesh拓?fù)渖蠝?zhǔn)確建模APR和拓?fù)涓兄贤ㄐ诺男袨?，并使用?nèi)部高精度仿真基礎(chǔ)設(shè)施校準(zhǔn)模型。

6.評估

在本節(jié)中，我們探索了UB-Mesh的架構(gòu)設(shè)計(jì)空間，并分析其相較于基線Clos架構(gòu)的優(yōu)勢。為便于更詳細(xì)的比較，我們將分析分為兩個層級：機(jī)架內(nèi)架構(gòu)比較和機(jī)架間架構(gòu)比較。

6.1 實(shí)驗(yàn)設(shè)置

表5：基準(zhǔn)模型

表5列出了基準(zhǔn)工作負(fù)載：Llama-70B和GPT3-175B為密集模型，而GPT4-2T為采用MoE技術(shù)的稀疏模型[20]。需要注意的是，GPT4的架構(gòu)尚未正式發(fā)布，我們采用了推測參數(shù)[23]。為評估系統(tǒng)在更大模型上的性能，我們還包括了Dense-1T和MoE-10T模型。為探索UB-Mesh的架構(gòu)，我們構(gòu)建了一個內(nèi)部仿真基礎(chǔ)設(shè)施，并與真實(shí)PoC硬件對齊，以評估集群規(guī)模的LLM訓(xùn)練性能。

6.2 機(jī)架內(nèi)架構(gòu)探索

圖16：基線機(jī)架內(nèi)架構(gòu)

我們比較了不同的機(jī)架內(nèi)網(wǎng)絡(luò)架構(gòu)，如圖16所示：

(a) 2D-FM：UB-Mesh的架構(gòu)。64個NPU通過電線直接互連，形成2D全互連網(wǎng)絡(luò)拓?fù)洌╔-Fullmesh + Y-Fullmesh）。該設(shè)計(jì)成本低，采用48個LRS（見表3）聚合機(jī)架間帶寬并實(shí)現(xiàn)CPU-NPU互連（為簡化起見，圖中省略了CPU）。

(b) 1D-FM-A：此替代架構(gòu)保留1D X-Fullmesh，即每個板上的8個NPU仍直接互連。然而，跨板通信通過36個LRS實(shí)現(xiàn)。每個NPU有一個UB x16 IO連接到LRS，另一個UB x16 IO連接到四個高基數(shù)交換機(jī)（HRS）用于機(jī)架間通信。

(c) 1D-FM-B：此架構(gòu)進(jìn)一步用HRS替換LRS。每個背板中的四個LRS用于NPU到CPU通信?？绨錘PU通信通過四個背板中的八個交換機(jī)實(shí)現(xiàn)。這些交換機(jī)還連接到機(jī)架間網(wǎng)絡(luò)，為每個NPU提供UB x32 IO用于機(jī)架間通信。

(d) Clos：此架構(gòu)不使用NPU間直接連接，而是將64個NPU的所有端口連接到72個LRS，形成對稱的Clos拓?fù)?。該架?gòu)提供最高靈活性，但需要大量交換資源。

圖17：不同機(jī)架內(nèi)拓?fù)涞男阅?/p>

性能比較：在圖17中，我們固定機(jī)架間架構(gòu)（2D-FM，見第6.3節(jié)），并比較不同機(jī)架內(nèi)架構(gòu)的訓(xùn)練吞吐量，相對于Clos基線。SuperPod規(guī)模為8000（128個機(jī)架），評估的序列長度范圍為8000至1000萬。我們計(jì)算了不同序列長度的平均性能。

如圖17-(a)所示，與Clos架構(gòu)相比，2D-FM架構(gòu)實(shí)現(xiàn)了93.2%至95.9%的訓(xùn)練性能。1D-FM-A架構(gòu)顯示出較低的性能下降，對于Llama2-70B相比2D-FM提升了2.44%的性能。對于其他參數(shù)更多的模型，提升幅度小于1.6%。1D-FM-B架構(gòu)由于更高的機(jī)架間帶寬，相比2D-FM顯示出略高的性能提升，超過3%，但提升幅度仍較小。

圖17-(b)探索了不同序列長度下的性能，相對于Clos架構(gòu)的基線場景（所有模型平均值）。對于8000至3.2萬的序列長度，2D-FM架構(gòu)實(shí)現(xiàn)了95.5%的性能，略低于1D-FM-A（98.1%）和1D-FM-B（99.2%）。對于6.4萬至1000萬的序列長度，2D-FM架構(gòu)相比Clos架構(gòu)實(shí)現(xiàn)了95.0%的性能。

可以看出，與Clos相比，2D-FM架構(gòu)以遠(yuǎn)低于硬件成本提供了相似的訓(xùn)練性能（性能差距在7%以內(nèi)），這將在第6.4節(jié)中進(jìn)行評估。

6.3 機(jī)架間架構(gòu)探索

圖18：機(jī)架間網(wǎng)絡(luò)架構(gòu)

在機(jī)架內(nèi)2D-FM之外，UB-Mesh在UB-Mesh Pod內(nèi)采用另一個2D-FM，形成4D-FM網(wǎng)絡(luò)架構(gòu)。如圖18所示，我們比較了UB-Mesh的2D全互連機(jī)架間架構(gòu)與基線Clos架構(gòu)：

(a) 2D-FM：16個機(jī)架在水平和垂直方向直接互連，形成2D全互連直接連接。所有機(jī)架連接到HRS交換機(jī)，支持跨Pod互連。在UB-Mesh-Pod中，可支持三種路由策略：  

• 最短路徑（Shortest）：基線路由策略，僅選擇2D網(wǎng)格上的最短路徑用于P2P通信。
• 繞行（Detour）：根據(jù)全路徑路由機(jī)制，還使用其他路徑以最大化帶寬。
• 借用（Borrow）：由于所有機(jī)架也連接到交換機(jī)，我們允許機(jī)架“借用”交換機(jī)的帶寬。

(b) Clos：所有機(jī)架連接到HRS交換機(jī)，無直接連接。此架構(gòu)比2D-FM消耗更多交換機(jī)，但提供最高的All-to-All帶寬和靈活性。

圖19：不同機(jī)架間互連的端到端性能比較

性能比較：如圖19所示，我們比較了2D-FM（包括不同路由策略）和Clos架構(gòu)的性能。可以看出，與理想的Clos架構(gòu)相比，2D-FM與Clos之間的性能差距微乎其微，尤其是應(yīng)用Detour和Borrow策略時。GPT3-175B對機(jī)架間通信性能不敏感，而在GPT4-2T上，采用最短路徑路由的樸素2D-FM訓(xùn)練性能下降0.73%。采用Detour和Borrow路由策略后，性能差距縮小至僅0.46%。總體而言，2D-FM機(jī)架間互連表現(xiàn)出與昂貴的Clos架構(gòu)幾乎相同的性能。

圖20：機(jī)架間帶寬探索

機(jī)架間帶寬探索：圖20-(b)比較了8000 SuperPod在不同機(jī)架間帶寬條件下的吞吐性能。評估的機(jī)架間帶寬為每NPU的x4、x8、x16和x32 UB IO。對于8000至3.2萬的序列長度，最優(yōu)機(jī)架間帶寬為UB x16；而對于6.4萬至1000萬的序列長度，最優(yōu)機(jī)架間帶寬為UB x32。從UB x8增加到UB x16在8000-3.2萬序列長度范圍內(nèi)的性能增益僅為0.44%，但從UB x16增加到UB x32在6.4萬-1000萬序列長度范圍內(nèi)的性能增益更為顯著，達(dá)1.85%。在6.4萬至1000萬序列長度的場景中，部分TP（張量并行）和SP（序列并行）流量不可避免地穿越機(jī)架間鏈路。更高的機(jī)架間帶寬顯著減少了TP和SP的通信時間，在這些場景中帶來更明顯的性能提升。這些數(shù)據(jù)強(qiáng)調(diào)了將機(jī)架間帶寬與不同模型場景的特定序列長度需求相匹配的重要性，特別是在減少大規(guī)模模型的TP和SP通信延遲方面。

UB-Mesh默認(rèn)分配每NPU UB x16 IO用于機(jī)架間通信，以在成本和性能之間實(shí)現(xiàn)平衡。我們還可調(diào)整機(jī)架內(nèi)/機(jī)架間帶寬比例，以匹配特定LLM訓(xùn)練工作負(fù)載的需求。

6.4 成本效益比較

圖21：CapEx比較

系統(tǒng)成本通常以總擁有成本（TCO）衡量，即TCO = CapEx + OpEx?？紤]到相對于基線的訓(xùn)練性能，我們定義系統(tǒng)的成本效益如下：

我們使用內(nèi)部數(shù)據(jù)估算CapEx成本，包括NPU、CPU、LRS、HRS、電纜和其他模塊的成本，并比較不同架構(gòu)。如圖21所示，UB-Mesh的4D-FM+Clos架構(gòu)相較于2D-FM+x16（表示每NPU UB x16 IO）Clos、1D-FM+x16 Clos和x64T Clos架構(gòu)分別實(shí)現(xiàn)了1.18倍、1.26倍、1.65倍和2.46倍的CapEx降低。與基線Clos架構(gòu)相比，UB-Mesh成功將系統(tǒng)中網(wǎng)絡(luò)基礎(chǔ)設(shè)施成本占比從67%降低至20%，這得益于節(jié)省了高性能交換機(jī)和長距離光纜/模塊。根據(jù)我們的評估，與基線Clos架構(gòu)相比，98%的高基數(shù)交換機(jī)和93%的光模塊得以節(jié)省。

OpEx降低主要包括系統(tǒng)生命周期內(nèi)的電費(fèi)和維護(hù)成本。UB-Mesh相較于Clos減少了約35%的OpEx，這得益于其大幅減少了交換機(jī)和光模塊的使用。根據(jù)我們云部門的AI系統(tǒng)估算，OpEx約占TCO的30%。最終，根據(jù)公式(1)，UB-Mesh實(shí)現(xiàn)了2.04倍更高的成本效益。

6.5 線性度評估

圖22：序列長度256K時的線性度分析

AI集群的線性度指的是集群性能隨NPU數(shù)量增加的線性擴(kuò)展程度。具體而言，我們通過以下公式測量線性度：

圖22評估了UB-Mesh在不同集群規(guī)模下的線性度?；鶞?zhǔn)規(guī)模（即圖中的1×規(guī)模）因任務(wù)而異。具體而言，Llama2-70B使用128個NPU，GPT3-175B的基準(zhǔn)規(guī)模為512，Dense-1T和GPT4-2T使用1000個NPU。

可以看出，UB-Mesh在所有任務(wù)的1×至32×規(guī)模下線性度超過100%，這是因?yàn)橐?guī)模提升提供了更多高帶寬域，并釋放了搜索更優(yōu)并行策略以提升MFU（模型浮點(diǎn)利用率）的潛力。GPT4-2T和Dense-1T模型在64×規(guī)模（涉及6.4萬個NPU）下線性度有所下降，但仍高于95%。

6.6 網(wǎng)絡(luò)可靠性分析

AI訓(xùn)練過程對硬件故障敏感。經(jīng)典的基于時間的指標(biāo)用于評估系統(tǒng)可靠性：

其中，MTBF表示平均故障間隔時間，MTTR表示平均修復(fù)時間。

表6：MTBF估算

如表6估算，優(yōu)先使用直接連接的電線（E-Cables）而非光纖和交換機(jī)，大幅降低了網(wǎng)絡(luò)模塊的年化故障率（AFR）。

然后，我們計(jì)算兩種架構(gòu)的MTBF

：基線Clos在8000個NPU集群中為13.8小時，而UB-Mesh達(dá)到98.5小時，提升了7.14倍。最終，根據(jù)公式(3)（我們根據(jù)現(xiàn)有統(tǒng)計(jì)數(shù)據(jù)假設(shè)MTTR為75分鐘），UB-Mesh的可用性為98.8%，顯著優(yōu)于Clos的91.6%（提升7.2%）。

為進(jìn)一步提升可用性，我們還精心設(shè)計(jì)了內(nèi)部網(wǎng)絡(luò)監(jiān)控工具，可在10分鐘內(nèi)快速識別和定位網(wǎng)絡(luò)故障，并在3分鐘內(nèi)觸發(fā)任務(wù)遷移，大幅減少M(fèi)TTR。通過這種MTTR優(yōu)化，根據(jù)我們的評估，UB-Mesh的可用性可進(jìn)一步提升至99.78%。

總結(jié)：與基線Clos架構(gòu)相比，UB-Mesh以微小的性能下降（7%以內(nèi)）為代價實(shí)現(xiàn)了2.04倍更高的成本效益。由于大幅減少了交換機(jī)和光模塊的使用，UB-Mesh將網(wǎng)絡(luò)可用性提升了7.2%。UB-Mesh還在多個LLM訓(xùn)練任務(wù)上實(shí)現(xiàn)了95%以上的線性度。

7.討論

集合通信協(xié)處理器

值得一提的是，在UB IO控制器中還配備了一個特殊的協(xié)處理器，稱為集合通信單元（Collective Communication Unit, CCU），用于卸載集合通信任務(wù)。具體而言，CCU執(zhí)行指令，主動從HBM讀取或?qū)懭霐?shù)據(jù)，啟動NPU間傳輸，并利用片上SRAM緩沖區(qū)執(zhí)行在線數(shù)據(jù)歸約操作。這種設(shè)計(jì)消除了從應(yīng)用程序內(nèi)存緩沖區(qū)到通信緩沖區(qū)的冗余數(shù)據(jù)拷貝，有效緩解了HBM帶寬的消耗，并通過基于檢查位的細(xì)粒度同步機(jī)制維持確定的歸約順序。CCU還能與計(jì)算核心無縫協(xié)作，實(shí)現(xiàn)高效的計(jì)算與通信重疊。

支持大規(guī)模專家模型

除了密集模型和常規(guī)MoE模型外，我們注意到人們正在積極探索包含大規(guī)模專家的MoE模型[6, 7]。這種設(shè)計(jì)通常需要大規(guī)模、細(xì)粒度的All-to-All通信。UB-Mesh通過多路徑和層次化的All-to-All優(yōu)化（見第5.1節(jié)）以及基于UB的Load/Store數(shù)據(jù)傳輸，能夠高效支持此類大規(guī)模專家模型。CCU還能有效卸載All-to-All操作，節(jié)省寶貴的計(jì)算核心使用，這與[7]中的預(yù)期一致。

8.結(jié)論

本文提出了UB-Mesh，一種為下一代大規(guī)模語言模型（LLM）訓(xùn)練設(shè)計(jì)的新型數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)。UB-Mesh采用n維全互連（nD-FullMesh）網(wǎng)絡(luò)拓?fù)洌瑴p少了交換機(jī)和光模塊的使用，同時適配LLM工作負(fù)載的分層局部化流量模式。我們深入考慮了架構(gòu)設(shè)計(jì)、物理實(shí)現(xiàn)和網(wǎng)絡(luò)系統(tǒng)優(yōu)化，并提出了多項(xiàng)技術(shù)以應(yīng)對各種挑戰(zhàn)。與傳統(tǒng)的Clos網(wǎng)絡(luò)相比，UB-Mesh在提供相似LLM訓(xùn)練性能的同時，實(shí)現(xiàn)了2.04倍的系統(tǒng)級成本效益提升和7.2%的可用性提升。此外，UB-Mesh在多項(xiàng)LLM訓(xùn)練任務(wù)上實(shí)現(xiàn)了95%以上的線性擴(kuò)展能力。

本文轉(zhuǎn)載自??Andy730??，作者：常華Andy